Files
YG_FT/backend/app/api/v1/endpoints/data_process.py
2026-07-27 14:41:38 +08:00

1581 lines
58 KiB
Python

from __future__ import annotations
import hashlib
import ipaddress
import json
import logging
import os
import re
import socket
from contextlib import contextmanager
from dataclasses import asdict
from pathlib import Path
from typing import Any, Iterator, Literal
from urllib.parse import quote, urlsplit
import psycopg
from fastapi import (
APIRouter,
BackgroundTasks,
Body,
Depends,
File,
Header,
HTTPException,
Query,
UploadFile,
)
from fastapi.responses import StreamingResponse
from psycopg.rows import dict_row
from app.modules.data_process.algorithms import (
ParsedText,
canonical_record_json,
content_quality_flags,
desensitize_pii,
desensitize_structured_record,
detect_pdf_document_noise,
estimate_token_count,
extract_pdf_page_texts,
generate_standard_records,
is_near_duplicate,
near_duplicate_fingerprint,
parse_text_content,
preprocess_structured_records,
remove_document_noise,
score_quality,
)
from app.modules.data_process.document_chunking import (
DocumentChunk,
chunk_fixed_text,
chunk_layout_document,
chunk_semantic_text,
merge_short_chunks,
)
from app.modules.data_process.generation import generate_model_records
from app.modules.data_process.storage import (
LocalDataProcessStorage,
StagedSourceObject,
get_data_process_storage,
)
from app.modules.data_process.store import (
ConflictError,
DataProcessStore,
DataProcessStoreError,
InvalidStateError,
NotFoundError,
get_data_process_store,
new_id,
)
from app.schemas.data_process import (
DataProcessRegenerateRequest,
DataProcessStatus,
DataProcessTaskCreate,
DataProcessTaskUpdate,
ExternalPullRequest,
ExternalSourceRequest,
GenerateRequest,
PreviewBuildRequest,
PreviewItemCreate,
PreviewItemUpdate,
ProcessType,
PublishRequest,
ResultUpdate,
)
router = APIRouter(prefix="/data-process")
logger = logging.getLogger(__name__)
MAX_SOURCE_FILE_BYTES = 200 * 1024 * 1024
MAX_SOURCE_FILE_COUNT = 20
MAX_SOURCE_BATCH_BYTES = 500 * 1024 * 1024
MAX_EXTERNAL_PULL_BYTES = 50 * 1024 * 1024
STRUCTURED_SOURCE_SUFFIXES = frozenset(
{".json", ".jsonl", ".ndjson", ".csv", ".tsv", ".xlsx"}
)
UNSTRUCTURED_SOURCE_SUFFIXES = frozenset(
{
".txt",
".md",
".markdown",
".pdf",
".docx",
".pptx",
".json",
".jsonl",
".ndjson",
}
)
SUPPORTED_SOURCE_SUFFIXES = STRUCTURED_SOURCE_SUFFIXES | UNSTRUCTURED_SOURCE_SUFFIXES
LEGACY_OFFICE_CONVERSIONS = {
".doc": ".docx",
".xls": ".xlsx",
".ppt": ".pptx",
}
def ok(data: Any = None, message: str = "ok") -> dict[str, Any]:
return {"code": 0, "message": message, "data": data}
def fail(status_code: int, message: str) -> HTTPException:
return HTTPException(
status_code=status_code,
detail={"code": status_code, "message": message, "data": None},
)
@contextmanager
def api_errors() -> Iterator[None]:
try:
yield
except NotFoundError as exc:
raise fail(404, str(exc)) from exc
except ConflictError as exc:
raise fail(409, str(exc)) from exc
except InvalidStateError as exc:
raise fail(409, str(exc)) from exc
except (DataProcessStoreError, ValueError) as exc:
raise fail(400, str(exc)) from exc
except psycopg.errors.UndefinedTable as exc:
raise fail(503, "data process schema is not installed; run schema_cli --check") from exc
except psycopg.OperationalError as exc:
raise fail(503, "data process database is unavailable") from exc
def _safe_file_name(value: str | None, fallback: str) -> str:
name = Path((value or "").replace("\\", "/")).name.replace("\x00", "").strip()
return name if name not in {"", ".", ".."} else fallback
def _range_not_satisfiable(size: int) -> HTTPException:
return HTTPException(
status_code=416,
detail={"code": 416, "message": "invalid source byte range", "data": None},
headers={"Content-Range": f"bytes */{size}"},
)
def _source_byte_range(value: str | None, size: int) -> tuple[int, int] | None:
if value is None:
return None
match = re.fullmatch(r"bytes=(\d*)-(\d*)", value.strip())
if match is None or size <= 0:
raise _range_not_satisfiable(size)
start_text, end_text = match.groups()
if not start_text and not end_text:
raise _range_not_satisfiable(size)
if start_text:
start = int(start_text)
end = int(end_text) if end_text else size - 1
if start >= size or end < start:
raise _range_not_satisfiable(size)
else:
suffix_length = int(end_text)
if suffix_length <= 0:
raise _range_not_satisfiable(size)
start = max(0, size - suffix_length)
end = size - 1
return start, min(end, size - 1)
def _commit_source_batch(
store: DataProcessStore,
storage: LocalDataProcessStorage,
task_id: str,
prepared: list[dict[str, Any]],
staged: list[StagedSourceObject],
) -> list[dict[str, Any]]:
storage.publish(staged)
try:
return store.add_source_files(task_id, prepared)
except Exception:
for item in staged:
try:
storage.delete(item.reference)
except Exception:
# 文件系统回滚失败不能覆盖数据库抛出的根因,并继续清理其余对象。
pass
raise
def _value(config: dict[str, Any], snake_name: str, camel_name: str, default: Any) -> Any:
if snake_name in config:
return config[snake_name]
return config.get(camel_name, default)
def _preprocess_options(config: dict[str, Any]) -> set[str]:
values = _value(config, "preprocess_options", "preprocessOptions", [])
return {str(item) for item in values} if isinstance(values, list) else set()
def _preview_quality(content: str, config: dict[str, Any]) -> dict[str, Any]:
records = generate_standard_records(
[{"id": "quality-preview", "edited_content": content}],
split={"train": 100, "validation": 0, "test": 0},
)
record = records[0] if records else {"instruction": "", "input": "", "output": ""}
minimum = int(_value(config, "min_output_length", "minOutputLength", 20) or 20)
return asdict(
score_quality(
record,
min_output_length=max(1, minimum),
source_content=content,
)
)
def _parse_stored_source(source: dict[str, Any]) -> ParsedText:
"""重新读取已入库的规范化正文,避免把二进制格式当二进制重复解析。"""
content = str(source.get("content") or "")
file_format = str(source.get("file_format") or "").lower()
if file_format == "xlsx":
# XLSX 上传阶段已安全解析为 JSONL 后入库。
return parse_text_content(content, file_format="jsonl")
if file_format in {"pdf", "docx", "pptx"}:
# 文档上传阶段已抽取文本,预览阶段只需要对正文切片。
return ParsedText(format=file_format, text=content, records=())
return parse_text_content(
content,
filename=source.get("name"),
file_format=file_format or None,
)
def _chunk_source_text(
source: dict[str, Any],
config: dict[str, Any],
preprocess_options: set[str],
) -> list[DocumentChunk]:
"""根据任务配置调用真实的 Docling/LlamaIndex 切分器。"""
method = str(_value(config, "chunk_method", "chunkMethod", "layout_hybrid"))
preserve_context = "preserve_context" in preprocess_options
chunk_size = int(_value(config, "chunk_size", "chunkSize", 800))
overlap = (
int(_value(config, "chunk_overlap", "chunkOverlap", 100))
if preserve_context
else 0
)
text = str(source.get("content") or "")
if method == "layout_hybrid":
raw = source.get("raw_content")
if not isinstance(raw, bytes):
raise InvalidStateError("版面结构混合切分需要原始文件,请重新上传后再处理")
chunks = chunk_layout_document(
raw,
filename=str(source.get("name") or "document.pdf"),
source_text=text,
chunk_size=chunk_size,
)
elif method == "semantic":
chunks = chunk_semantic_text(
text,
chunk_size=chunk_size,
chunk_overlap=overlap,
breakpoint_percentile_threshold=int(
_value(
config,
"semantic_breakpoint_percentile",
"semanticBreakpointPercentile",
95,
)
),
)
elif method == "fixed":
chunks = chunk_fixed_text(text, chunk_size=chunk_size, chunk_overlap=overlap)
else:
raise ValueError(f"unsupported chunk method: {method}")
if "merge_short_content" in preprocess_options:
chunks = merge_short_chunks(
chunks,
source_text=text,
min_token_count=int(_value(config, "min_chunk_size", "minChunkSize", 100)),
max_token_count=chunk_size,
)
return chunks
_NEGATION_MARKERS = frozenset({"", "", "", "", "没有", "并非", "not", "no", "never"})
def _safe_near_duplicate(left: str, right: str) -> bool:
"""保守判断近重复,数字或否定含义变化时始终保留。"""
if min(estimate_token_count(left), estimate_token_count(right)) < 20:
return False
if re.findall(r"\d+(?:\.\d+)?", left) != re.findall(r"\d+(?:\.\d+)?", right):
return False
left_lower = left.casefold()
right_lower = right.casefold()
left_negations = {marker for marker in _NEGATION_MARKERS if marker in left_lower}
right_negations = {marker for marker in _NEGATION_MARKERS if marker in right_lower}
if left_negations != right_negations:
return False
return is_near_duplicate(
left,
right,
similarity_threshold=0.92,
max_hamming_distance=2,
)
def _near_duplicate_band_keys(content: str) -> tuple[tuple[int, int], ...]:
"""将 64 位 SimHash 分为三段,汉明距离不超过 2 时至少命中一段。"""
fingerprint = int(near_duplicate_fingerprint(content), 16)
widths = (22, 21, 21)
shift = 0
keys: list[tuple[int, int]] = []
for index, width in enumerate(widths):
keys.append((index, (fingerprint >> shift) & ((1 << width) - 1)))
shift += width
return tuple(keys)
def _build_preview_items(
task: dict[str, Any], source_files: list[dict[str, Any]]
) -> list[dict[str, Any]]:
config = task.get("config") or {}
process_type = task["process_type"]
preprocess_options = _preprocess_options(config)
should_desensitize = "desensitize" in preprocess_options
should_clean_invalid = bool(
preprocess_options & {"clean_invalid", "clean_invalid_content"}
)
should_deduplicate = bool(
preprocess_options & {"deduplicate", "deduplicate_content"}
)
seen_content_hashes: set[str] = set()
seen_near_duplicate_bands: dict[tuple[int, int], list[str]] = {}
items: list[dict[str, Any]] = []
def append_item(item: dict[str, Any]) -> None:
content = str(item.get("edited_content") or "").strip()
if should_clean_invalid and not content:
return
content_hash = hashlib.sha256(content.encode("utf-8")).hexdigest()
if should_deduplicate and content_hash in seen_content_hashes:
return
seen_content_hashes.add(content_hash)
if not content:
item["status"] = "invalid"
items.append(item)
for source in source_files:
parsed = _parse_stored_source(source)
if process_type == "unstructured":
document_noise_spans = tuple(source.get("document_noise_spans") or ())
chunks = _chunk_source_text(source, config, preprocess_options)
for chunk in chunks:
content = (
remove_document_noise(
chunk.contextualized_content,
document_noise_spans,
source_offset=chunk.source_start or 0,
)
if (
should_clean_invalid
and document_noise_spans
and chunk.source_start is not None
)
else chunk.contextualized_content
)
preprocess_flags = content_quality_flags(
content,
min_chars=0,
min_tokens=0,
)
if content != chunk.contextualized_content:
preprocess_flags = (*preprocess_flags, "document_noise_removed")
flag_set = set(preprocess_flags)
if "clean_invalid_content" in preprocess_options and flag_set & {
"empty_content",
"low_printable_ratio",
"repetitive_content",
}:
continue
if "filter_low_quality" in preprocess_options and flag_set & {
"content_too_long",
"mojibake",
"low_printable_ratio",
"repetitive_content",
}:
continue
if "deduplicate_content" in preprocess_options:
band_keys = _near_duplicate_band_keys(content)
candidates = {
previous
for key in band_keys
for previous in seen_near_duplicate_bands.get(key, ())
}
if any(
_safe_near_duplicate(content, previous)
for previous in candidates
):
continue
for key in band_keys:
seen_near_duplicate_bands.setdefault(key, []).append(content)
pii_counts: dict[str, int] = {}
if should_desensitize:
content, pii_counts = desensitize_pii(content)
quality = _preview_quality(content, config)
quality["pii_replacements"] = pii_counts
quality["preprocess_flags"] = list(preprocess_flags)
quality["chunk_method"] = str(
_value(config, "chunk_method", "chunkMethod", "layout_hybrid")
)
quality["heading_path"] = list(chunk.heading_path)
quality["source_pages"] = list(chunk.source_pages)
quality["doc_item_refs"] = list(chunk.doc_item_refs)
quality["source_bboxes"] = list(chunk.source_bboxes)
append_item(
{
"source_file_id": source["id"],
"original_content": chunk.original_content,
"edited_content": content,
"source_start": chunk.source_start,
"source_end": chunk.source_end,
"source_start_line": chunk.source_start_line,
"source_end_line": chunk.source_end_line,
"token_count": estimate_token_count(content),
"status": (
"modified"
if content != chunk.original_content
else "original"
),
"quality_score": quality,
}
)
continue
structured_options = preprocess_options & {
"clean_invalid",
"detect_structure",
"deduplicate",
"normalize_format",
"filter_anomaly",
}
source_records = list(parsed.records)
processed_records = preprocess_structured_records(
source_records,
structured_options,
)
if not processed_records and parsed.text and not source_records:
processed_records = [{"value": parsed.text}]
same_cardinality = len(processed_records) == len(source_records)
for index, record in enumerate(processed_records):
original_record = source_records[index] if same_cardinality else record
original_content = json.dumps(
original_record,
ensure_ascii=False,
separators=(",", ":"),
)
pii_counts: dict[str, int] = {}
edited_record = record
if should_desensitize:
edited_record, pii_counts = desensitize_structured_record(record)
content = (
canonical_record_json(edited_record)
if "normalize_format" in preprocess_options
else json.dumps(
edited_record,
ensure_ascii=False,
separators=(",", ":"),
)
)
quality = _preview_quality(content, config)
quality["pii_replacements"] = pii_counts
append_item(
{
"source_file_id": source["id"],
"original_content": original_content,
"edited_content": content,
"source_start": None,
"source_end": None,
"source_start_line": None,
"source_end_line": None,
"token_count": estimate_token_count(content),
"status": "modified" if content != original_content else "original",
"quality_score": quality,
}
)
return items
def _all_preview_items(store: DataProcessStore, task_id: str) -> list[dict[str, Any]]:
"""分页读取全部预览项,避免固定上限静默截断任务。"""
items: list[dict[str, Any]] = []
page = 1
page_size = 5_000
while True:
result = store.list_preview_items(task_id, page=page, page_size=page_size)
batch = result["items"]
items.extend(batch)
if len(items) >= int(result["total"]) or not batch:
return items
page += 1
def _run_generation(
store: DataProcessStore, task_id: str, generation_run_id: str
) -> None:
try:
task = store.get_task(task_id)
if not store.generation_is_running(task_id, generation_run_id):
return
all_preview_items = _all_preview_items(store, task_id)
preview_items = [
item
for item in all_preview_items
if item.get("status") != "invalid"
and str(item.get("edited_content") or item.get("original_content") or "").strip()
]
pre_filtered_count = len(all_preview_items) - len(preview_items)
config = task.get("config") or {}
model_id = _value(config, "generation_model_id", "generationModelId", None)
generation_model: dict[str, Any] | None = None
if model_id:
generation_model = store.get_generation_model(str(model_id))
task = store.save_generation_model_snapshot(
task_id,
generation_model,
generation_run_id=generation_run_id,
)
config = task.get("config") or config
split = _value(
config,
"dataset_split",
"datasetSplit",
{"train": 80, "validation": 10, "test": 10},
)
pairs = (
_value(config, "qa_pairs_per_chunk", "qaPairsPerChunk", 1)
if task["process_type"] == "unstructured"
else _value(config, "qa_pairs_per_row", "qaPairsPerRow", 1)
)
output_type = str(
_value(config, "output_type", "outputType", "standard")
).strip().lower()
if generation_model:
runtime_config = {
**config,
"generation_prompt": _value(
config, "generation_prompt", "generationPrompt", ""
),
"output_type": output_type,
"reasoning_detail": _value(
config, "reasoning_detail", "reasoningDetail", "normal"
),
"max_tokens": _value(config, "max_tokens", "maxTokens", 1024),
"json_mode": _value(config, "json_mode", "jsonMode", False),
}
def report_progress(processed_count: int, total_count: int) -> None:
if not store.update_generation_progress(
task_id,
generation_run_id,
processed_count,
total_count,
):
raise InvalidStateError("generation run is no longer active")
generated = generate_model_records(
preview_items,
model=generation_model,
config=runtime_config,
task_id=task_id,
split=split,
qa_pairs_per_item=int(pairs or 1),
on_progress=report_progress,
)
elif output_type == "reasoning":
raise InvalidStateError("思维链输出必须配置可用的数据生成模型")
else:
generated = generate_standard_records(
preview_items,
qa_pairs_per_item=int(pairs or 1),
semantic_enrichment=bool(
_value(config, "semantic_enrichment", "semanticEnrichment", False)
),
split=split,
split_seed=task_id,
)
if not store.update_generation_progress(
task_id,
generation_run_id,
len(preview_items),
len(preview_items),
):
return
known_fingerprints: set[str] = set()
accepted: list[dict[str, Any]] = []
filtered_count = pre_filtered_count
duplicate_count = 0
error_count = 0
quality_filter = bool(
_value(config, "quality_filter_enabled", "qualityFilterEnabled", False)
)
filter_low = bool(_value(config, "filter_low_quality", "filterLowQuality", True))
filter_short = bool(
_value(config, "filter_short_content", "filterShortContent", True)
)
deduplicate = bool(
_preprocess_options(config)
& {"deduplicate", "deduplicate_content"}
)
minimum = max(1, int(_value(config, "min_output_length", "minOutputLength", 20) or 20))
preview_sources = {
str(item["id"]): str(
item.get("edited_content") or item.get("original_content") or ""
)
for item in preview_items
}
for record in generated:
quality = score_quality(
record,
min_output_length=minimum,
source_content=preview_sources.get(str(record.get("preview_item_id") or ""), ""),
known_fingerprints=known_fingerprints,
)
if "duplicate_record" in quality.flags:
duplicate_count += 1
else:
# 即使首条随后因短文本/低质量被过滤,也要阻止同批后续重复结果。
known_fingerprints.add(quality.fingerprint)
should_filter = (
(deduplicate and "duplicate_record" in quality.flags)
or (
quality_filter
and filter_short
and "output_too_short" in quality.flags
)
or (quality_filter and filter_low and not quality.is_valid)
)
if not quality.is_valid:
error_count += 1
record["status"] = "invalid"
record["error"] = ", ".join(quality.flags) or "quality validation failed"
if should_filter:
filtered_count += 1
continue
record["quality_score"] = asdict(quality)
accepted.append(record)
# stop 请求可能在纯函数计算期间到达,最终写入前再次检查状态。
if store.generation_is_running(task_id, generation_run_id):
store.complete_generation(
task_id,
accepted,
generation_run_id=generation_run_id,
filtered_count=filtered_count,
duplicate_count=duplicate_count,
error_count=error_count,
)
except Exception as exc: # noqa: BLE001 - background failures must be persisted
try:
if store.generation_is_running(task_id, generation_run_id):
store.mark_failed(
task_id,
str(exc),
generation_run_id=generation_run_id,
)
except Exception:
return
@router.get("")
def list_tasks(
page: int = Query(default=1, ge=1),
page_size: int = Query(default=20, ge=1, le=200),
keyword: str | None = Query(default=None),
status: DataProcessStatus | None = Query(default=None),
process_type: ProcessType | None = Query(default=None),
store: DataProcessStore = Depends(get_data_process_store),
) -> dict[str, Any]:
with api_errors():
return ok(
store.list_tasks(
page=page,
page_size=page_size,
keyword=keyword,
status=status,
process_type=process_type,
)
)
@router.post("")
def create_task(
payload: DataProcessTaskCreate,
store: DataProcessStore = Depends(get_data_process_store),
) -> dict[str, Any]:
with api_errors():
task = store.create_task(payload.model_dump(mode="json"))
return ok(task, "data process task created")
@router.get("/{task_id}")
def task_detail(
task_id: str,
store: DataProcessStore = Depends(get_data_process_store),
) -> dict[str, Any]:
with api_errors():
# 兼容旧版曾在第五步前清空结果的异常任务;严格特征匹配且幂等。
store.recover_legacy_aborted_regeneration(task_id)
task = store.get_task(task_id)
source_files = store.list_source_files(task_id)
task["source_files"] = source_files
task["source_file_count"] = len(source_files)
return ok(task)
@router.put("/{task_id}")
def update_task(
task_id: str,
payload: DataProcessTaskUpdate,
store: DataProcessStore = Depends(get_data_process_store),
) -> dict[str, Any]:
with api_errors():
return ok(
store.update_task(task_id, payload.model_dump(exclude_unset=True, mode="json")),
"data process task updated",
)
@router.post("/{task_id}/regenerate")
def prepare_regeneration(
task_id: str,
payload: DataProcessRegenerateRequest,
store: DataProcessStore = Depends(get_data_process_store),
) -> dict[str, Any]:
with api_errors():
return ok(
store.prepare_regeneration(task_id, payload.model_dump(mode="json")),
"data process task prepared for regeneration",
)
@router.delete("/{task_id}")
def delete_task(
task_id: str,
store: DataProcessStore = Depends(get_data_process_store),
) -> dict[str, Any]:
with api_errors():
store.delete_task(task_id)
return ok({"deleted": task_id}, "data process task deleted")
@router.get("/{task_id}/source-files")
def source_files(
task_id: str,
store: DataProcessStore = Depends(get_data_process_store),
) -> dict[str, Any]:
with api_errors():
return ok({"files": store.list_source_files(task_id)})
@router.post("/{task_id}/source-files")
async def upload_source_files(
task_id: str,
files: list[UploadFile] = File(...),
store: DataProcessStore = Depends(get_data_process_store),
storage: LocalDataProcessStorage = Depends(get_data_process_storage),
) -> dict[str, Any]:
if not files:
raise fail(400, "at least one source file is required")
if len(files) > MAX_SOURCE_FILE_COUNT:
raise fail(413, f"a source batch may contain at most {MAX_SOURCE_FILE_COUNT} files")
prepared: list[dict[str, Any]] = []
staged: list[StagedSourceObject] = []
batch_id = storage.new_batch_id()
batch_size = 0
commit_attempted = False
try:
with api_errors():
task = store.get_task(task_id)
process_type = str(task["process_type"])
if process_type == "external":
raise InvalidStateError(
"external tasks must import data through the external source endpoint"
)
allowed_suffixes = (
UNSTRUCTURED_SOURCE_SUFFIXES
if process_type == "unstructured"
else STRUCTURED_SOURCE_SUFFIXES
)
for upload in files:
raw = await upload.read(MAX_SOURCE_FILE_BYTES + 1)
if len(raw) > MAX_SOURCE_FILE_BYTES:
raise fail(413, f"source file exceeds {MAX_SOURCE_FILE_BYTES} bytes")
name = _safe_file_name(upload.filename, "source.txt")
suffix = Path(name).suffix.lower()
if suffix in LEGACY_OFFICE_CONVERSIONS:
replacement = LEGACY_OFFICE_CONVERSIONS[suffix]
raise fail(
415,
f"legacy {suffix} format is not supported; "
f"convert the file to {replacement} and upload again",
)
if suffix not in SUPPORTED_SOURCE_SUFFIXES:
raise fail(415, f"unsupported source file format: {suffix or 'none'}")
if suffix not in allowed_suffixes:
raise fail(
415,
f"{suffix} is not supported for {process_type} data processing",
)
parsed = parse_text_content(raw, filename=name)
if not parsed.text:
raise fail(400, f"source file is empty: {name}")
batch_size += len(raw)
if batch_size > MAX_SOURCE_BATCH_BYTES:
raise fail(413, f"source batch exceeds {MAX_SOURCE_BATCH_BYTES} bytes")
source_file_id = new_id("dpsf")
staged_object = storage.stage_bytes(
batch_id=batch_id,
task_id=task_id,
source_file_id=source_file_id,
version=1,
name=name,
content=raw,
)
staged.append(staged_object)
record_count = len(parsed.records) or (1 if parsed.text else 0)
prepared.append(
{
"id": source_file_id,
"storage_object_id": staged_object.reference,
"name": name,
"content": parsed.text,
"raw_size": len(raw),
"checksum_sha256": hashlib.sha256(raw).hexdigest(),
"file_format": parsed.format,
"record_count": record_count,
"metadata": {
"content_type": upload.content_type or "text/plain",
"original_size_bytes": len(raw),
"original_checksum_sha256": hashlib.sha256(raw).hexdigest(),
},
"created_by": None,
}
)
# publish 无论成功或失败都会消费并清理暂存对象,外层不能再次 discard。
commit_attempted = True
created = _commit_source_batch(store, storage, task_id, prepared, staged)
finally:
if not commit_attempted:
storage.discard(staged)
return ok({"files": created}, "source files uploaded")
@router.get("/{task_id}/source-files/{file_id}/content")
def source_file_content(
task_id: str,
file_id: str,
start_line: int | None = Query(default=None, ge=1),
line_count: int = Query(default=200, ge=1, le=10_000),
offset: int = Query(default=0, ge=0),
limit: int = Query(default=100_000, ge=1, le=1_000_000),
store: DataProcessStore = Depends(get_data_process_store),
) -> dict[str, Any]:
with api_errors():
if start_line is not None:
return ok(store.source_content_lines(task_id, file_id, start_line, line_count))
return ok(store.source_content_window(task_id, file_id, offset, limit))
@router.get("/{task_id}/source-files/{file_id}/raw")
def source_file_raw(
task_id: str,
file_id: str,
range_header: str | None = Header(default=None, alias="Range"),
store: DataProcessStore = Depends(get_data_process_store),
storage: LocalDataProcessStorage = Depends(get_data_process_storage),
) -> StreamingResponse:
with api_errors():
source = store.get_source_file(task_id, file_id, include_content=False)
if str(source.get("file_format") or "").lower() != "pdf":
raise fail(415, "raw inline preview is only available for PDF source files")
storage_object_id = str(source.get("storage_object_id") or "")
actual_size = storage.file_size(
storage_object_id,
expected_task_id=task_id,
expected_source_file_id=file_id,
)
if actual_size is None:
raise fail(410, "the original PDF is unavailable for this legacy source file")
expected_size = int(source.get("size_bytes") or 0)
if actual_size != expected_size:
raise ValueError("source object size does not match metadata")
selected_range = _source_byte_range(range_header, actual_size)
start, end = selected_range or (0, actual_size - 1)
length = end - start + 1
name = _safe_file_name(str(source.get("name") or "source.pdf"), "source.pdf")
headers = {
"Accept-Ranges": "bytes",
"Cache-Control": "private, no-store",
"Content-Disposition": f"inline; filename*=UTF-8''{quote(name, safe='')}",
"Content-Length": str(length),
"X-Accel-Buffering": "no",
"X-Content-Type-Options": "nosniff",
}
checksum = str(source.get("checksum_sha256") or "")
if checksum:
headers["ETag"] = f'"{checksum}"'
if selected_range is not None:
headers["Content-Range"] = f"bytes {start}-{end}/{actual_size}"
body = storage.iter_bytes(
storage_object_id,
expected_task_id=task_id,
expected_source_file_id=file_id,
expected_size=actual_size,
start=start,
length=length,
)
return StreamingResponse(
body,
status_code=206 if selected_range is not None else 200,
media_type="application/pdf",
headers=headers,
)
@router.get("/{task_id}/source-files/{file_id}/pdf-pages")
def source_file_pdf_pages(
task_id: str,
file_id: str,
store: DataProcessStore = Depends(get_data_process_store),
storage: LocalDataProcessStorage = Depends(get_data_process_storage),
) -> dict[str, Any]:
"""返回切片全文偏移对应的 PDF 物理页范围。"""
with api_errors():
source = store.get_source_file(task_id, file_id, include_content=False)
if str(source.get("file_format") or "").lower() != "pdf":
raise fail(415, "PDF page mapping is only available for PDF source files")
storage_object_id = str(source.get("storage_object_id") or "")
actual_size = storage.file_size(
storage_object_id,
expected_task_id=task_id,
expected_source_file_id=file_id,
)
if actual_size is None:
raise fail(410, "the original PDF is unavailable for this legacy source file")
expected_size = int(source.get("size_bytes") or 0)
if actual_size != expected_size:
raise ValueError("source object size does not match metadata")
raw = b"".join(
storage.iter_bytes(
storage_object_id,
expected_task_id=task_id,
expected_source_file_id=file_id,
expected_size=actual_size,
)
)
pages = extract_pdf_page_texts(raw)
return ok(
{
"page_count": len(pages),
"pages": [
{
"page_number": page.page_number,
"source_start": page.source_start,
"source_end": page.source_end,
}
for page in pages
],
}
)
@router.delete("/{task_id}/source-files/{file_id}")
def delete_source_file(
task_id: str,
file_id: str,
store: DataProcessStore = Depends(get_data_process_store),
storage: LocalDataProcessStorage = Depends(get_data_process_storage),
) -> dict[str, Any]:
with api_errors():
source = store.get_source_file(task_id, file_id, include_content=False)
storage_object_id = str(source.get("storage_object_id") or "")
storage.validate_owner(
storage_object_id,
expected_task_id=task_id,
expected_source_file_id=file_id,
)
store.delete_source_file(task_id, file_id)
cleanup_pending = False
try:
storage.delete(
storage_object_id,
expected_task_id=task_id,
expected_source_file_id=file_id,
)
except Exception:
# 数据库软删除已经提交,不能再向客户端返回可重试的失败;保留逻辑引用,
# 由后续存储清理任务重试物理删除。
cleanup_pending = True
logger.exception(
"failed to remove data process source object after soft deletion",
extra={"task_id": task_id, "source_file_id": file_id},
)
return ok(
{"deleted": file_id, "storage_cleanup_pending": cleanup_pending},
"source file removed",
)
def _external_postgres_connection(payload: ExternalSourceRequest) -> psycopg.Connection[Any]:
kind = payload.type.strip().lower()
parsed_url = urlsplit(payload.url)
scheme = parsed_url.scheme.lower()
if kind not in {"postgres", "postgresql"} or scheme not in {"postgres", "postgresql"}:
raise fail(501, f"external data source type is not supported: {payload.type}")
if parsed_url.username or parsed_url.password:
raise fail(400, "database credentials must use the account and password fields")
if payload.auth_mode not in {"none", "basic"}:
raise fail(400, "PostgreSQL supports only none or basic authentication")
if payload.auth_mode == "basic" and not payload.username:
raise fail(400, "database username is required for basic authentication")
hostname = parsed_url.hostname
if not hostname:
raise fail(400, "external PostgreSQL URL must include a hostname")
allow_private = os.getenv("DATA_PROCESS_ALLOW_PRIVATE_EXTERNAL_DB", "").lower() in {
"1",
"true",
"yes",
}
if not allow_private:
try:
addresses = {
item[4][0]
for item in socket.getaddrinfo(
hostname,
parsed_url.port or 5432,
type=socket.SOCK_STREAM,
)
}
except socket.gaierror as exc:
raise fail(400, "external PostgreSQL hostname cannot be resolved") from exc
if any(
(address := ipaddress.ip_address(value)).is_private
or address.is_loopback
or address.is_link_local
or address.is_reserved
or address.is_unspecified
for value in addresses
):
raise fail(
403,
"private or local database addresses are disabled; "
"set DATA_PROCESS_ALLOW_PRIVATE_EXTERNAL_DB=true only in a trusted deployment",
)
kwargs: dict[str, Any] = {
"connect_timeout": 5,
"row_factory": dict_row,
"application_name": "yg-ft-data-process-readonly",
"options": "-c default_transaction_read_only=on -c statement_timeout=30000",
}
if payload.auth_mode == "basic" and payload.username:
kwargs["user"] = payload.username
if payload.auth_mode == "basic" and payload.password:
kwargs["password"] = payload.password
return psycopg.connect(payload.url, **kwargs)
@router.post("/{task_id}/external/test")
def test_external_source(
task_id: str,
payload: ExternalSourceRequest,
store: DataProcessStore = Depends(get_data_process_store),
) -> dict[str, Any]:
with api_errors():
task = store.get_task(task_id)
if str(task.get("process_type")) != "external":
raise InvalidStateError(
"external source access requires an external data processing task"
)
try:
with _external_postgres_connection(payload) as conn:
conn.execute("SELECT 1 AS ok").fetchone()
except psycopg.Error as exc:
raise fail(502, "external PostgreSQL connection test failed") from exc
return ok({"connected": True, "type": payload.type})
@router.post("/{task_id}/external/pull")
def pull_external_source(
task_id: str,
payload: ExternalPullRequest,
store: DataProcessStore = Depends(get_data_process_store),
storage: LocalDataProcessStorage = Depends(get_data_process_storage),
) -> dict[str, Any]:
query = (payload.query or "").strip()
if query.endswith(";"):
query = query[:-1].rstrip()
if ";" in query:
raise fail(400, "external pull accepts exactly one read-only query")
first_token = query.split(maxsplit=1)[0].lower() if query else ""
if first_token not in {"select", "with"}:
raise fail(400, "a read-only SELECT or WITH query is required for external pull")
with api_errors():
task = store.get_task(task_id)
if str(task.get("process_type")) != "external":
raise InvalidStateError("external pull requires an external data processing task")
try:
with _external_postgres_connection(payload) as conn:
conn.execute("SET TRANSACTION READ ONLY")
conn.execute("SET LOCAL statement_timeout = '30s'")
cursor = conn.execute(query)
rows: list[dict[str, Any]] = []
content_parts: list[str] = []
content_size = 0
while len(rows) < payload.limit:
batch = cursor.fetchmany(min(1_000, payload.limit - len(rows)))
if not batch:
break
for row in batch:
line = json.dumps(row, ensure_ascii=False, default=str) + "\n"
content_size += len(line.encode("utf-8"))
if content_size > MAX_EXTERNAL_PULL_BYTES:
raise fail(413, "external pull result exceeds the 50 MiB safety limit")
rows.append(row)
content_parts.append(line)
conn.rollback()
except psycopg.Error as exc:
raise fail(502, "external PostgreSQL query failed") from exc
if not rows:
raise fail(400, "external query returned no rows")
content = "".join(content_parts)
raw = content.encode("utf-8")
name = _safe_file_name(payload.file_name, "external-data.jsonl")
source_file_id = new_id("dpsf")
staged = storage.stage_bytes(
batch_id=storage.new_batch_id(),
task_id=task_id,
source_file_id=source_file_id,
version=1,
name=name,
content=raw,
)
sources = _commit_source_batch(
store,
storage,
task_id,
[
{
"id": source_file_id,
"storage_object_id": staged.reference,
"name": name,
"content": content,
"raw_size": len(raw),
"checksum_sha256": hashlib.sha256(raw).hexdigest(),
"file_format": "jsonl",
"record_count": len(rows),
"metadata": {
"external_type": payload.type,
"external_host": urlsplit(payload.url).hostname,
"external_limit": payload.limit,
},
}
],
[staged],
)
source = sources[0]
return ok({"files": [source]}, "external source pulled")
def _prepare_preview_items(
task_id: str,
store: DataProcessStore,
storage: LocalDataProcessStorage,
source_file_ids: list[str] | None = None,
) -> list[dict[str, Any]]:
task = store.get_task(task_id)
source_summaries = store.list_source_files(task_id)
if source_file_ids is not None:
requested = set(source_file_ids)
source_summaries = [item for item in source_summaries if item["id"] in requested]
found = {item["id"] for item in source_summaries}
missing = requested - found
if missing:
raise NotFoundError(f"source files not found: {', '.join(sorted(missing))}")
sources = [
store.get_source_file(task_id, item["id"], include_content=True)
for item in source_summaries
]
if not sources:
raise InvalidStateError("at least one source file is required")
config = task.get("config") or {}
preprocess_options = _preprocess_options(config)
chunk_method = str(
_value(config, "chunk_method", "chunkMethod", "layout_hybrid")
)
is_unstructured = task.get("process_type") == "unstructured"
if is_unstructured and (
chunk_method == "layout_hybrid"
or preprocess_options & {"clean_invalid", "clean_invalid_content"}
):
for index, source in enumerate(sources):
if (
chunk_method != "layout_hybrid"
and str(source.get("file_format") or "").lower() != "pdf"
):
continue
storage_object_id = str(source.get("storage_object_id") or "")
actual_size = storage.file_size(
storage_object_id,
expected_task_id=task_id,
expected_source_file_id=str(source["id"]),
)
if actual_size is None:
if chunk_method == "layout_hybrid":
raise InvalidStateError(
"版面结构混合切分无法读取原始文件,请重新上传后再处理"
)
continue
expected_size = int(source.get("size_bytes") or 0)
if expected_size and actual_size != expected_size:
raise ValueError("source object size does not match metadata")
raw = b"".join(
storage.iter_bytes(
storage_object_id,
expected_task_id=task_id,
expected_source_file_id=str(source["id"]),
expected_size=actual_size,
)
)
enriched = dict(source)
if chunk_method == "layout_hybrid":
enriched["raw_content"] = raw
sources[index] = enriched
continue
if str(source.get("file_format") or "").lower() != "pdf":
continue
pages = extract_pdf_page_texts(raw)
extracted_text = "\n\n".join(page.text for page in pages if page.text)
if extracted_text != str(source.get("content") or ""):
logger.warning(
"skip PDF document noise detection because stored offsets differ for %s",
source["id"],
)
continue
enriched["document_noise_spans"] = detect_pdf_document_noise(pages)
sources[index] = enriched
items = _build_preview_items(task, sources)
if not items and source_file_ids is None:
raise InvalidStateError("source files did not produce preview items")
return items
@router.post("/{task_id}/preview/build")
def build_preview(
task_id: str,
payload: PreviewBuildRequest = Body(default_factory=PreviewBuildRequest),
store: DataProcessStore = Depends(get_data_process_store),
storage: LocalDataProcessStorage = Depends(get_data_process_storage),
) -> dict[str, Any]:
with api_errors():
selected_ids = payload.source_file_ids
items = _prepare_preview_items(task_id, store, storage, selected_ids)
created = store.replace_preview_items(
task_id,
items,
source_file_ids=selected_ids,
)
target_ids = selected_ids or [
str(source["id"]) for source in store.list_source_files(task_id)
]
file_counts = dict.fromkeys(target_ids, 0)
for item in created:
source_file_id = str(item.get("source_file_id") or "")
if source_file_id in file_counts:
file_counts[source_file_id] += 1
files = [
{
"source_file_id": source_file_id,
"preview_count": count,
"status": "completed" if count else "empty",
}
for source_file_id, count in file_counts.items()
]
return ok(
{
"items": created,
"total": len(created),
"page": 1,
"page_size": len(created),
"file_counts": file_counts,
"files": files,
},
"preview built",
)
@router.get("/{task_id}/preview")
def preview_items(
task_id: str,
source_file_id: str | None = Query(default=None),
page: int = Query(default=1, ge=1),
page_size: int = Query(default=200, ge=1, le=1000),
keyword: str | None = Query(default=None),
store: DataProcessStore = Depends(get_data_process_store),
) -> dict[str, Any]:
with api_errors():
return ok(
store.list_preview_items(
task_id,
source_file_id=source_file_id,
page=page,
page_size=page_size,
keyword=keyword,
)
)
@router.post("/{task_id}/preview")
def create_preview_item(
task_id: str,
payload: PreviewItemCreate,
store: DataProcessStore = Depends(get_data_process_store),
) -> dict[str, Any]:
with api_errors():
task = store.get_task(task_id)
item_payload = payload.model_dump(mode="json")
content = payload.edited_content
item_payload["token_count"] = estimate_token_count(content)
item_payload["quality_score"] = _preview_quality(
content,
task.get("config") or {},
)
if not content.strip():
item_payload["status"] = "invalid"
item = store.create_preview_item(task_id, item_payload)
return ok(item, "preview item created")
@router.put("/{task_id}/preview/{preview_id}")
def update_preview_item(
task_id: str,
preview_id: str,
payload: PreviewItemUpdate,
store: DataProcessStore = Depends(get_data_process_store),
) -> dict[str, Any]:
with api_errors():
task = store.get_task(task_id)
update = payload.model_dump(exclude_unset=True, mode="json")
update["quality_score"] = _preview_quality(payload.edited_content, task.get("config") or {})
item = store.update_preview_item(
task_id,
preview_id,
update,
)
return ok(item, "preview item updated")
@router.delete("/{task_id}/preview/{preview_id}")
def delete_preview_item(
task_id: str,
preview_id: str,
store: DataProcessStore = Depends(get_data_process_store),
) -> dict[str, Any]:
with api_errors():
store.delete_preview_item(task_id, preview_id)
return ok({"deleted": preview_id}, "preview item deleted")
def _start_generation(
task_id: str,
payload: GenerateRequest,
background_tasks: BackgroundTasks,
store: DataProcessStore,
) -> dict[str, Any]:
with api_errors():
task = store.start_generation(task_id, replace_existing=payload.replace_existing)
background_tasks.add_task(
_run_generation,
store,
task_id,
str(task["generation_run_id"]),
)
return ok(store.progress(task_id), "data process generation started")
@router.post("/{task_id}/generate")
def generate(
task_id: str,
background_tasks: BackgroundTasks,
payload: GenerateRequest = Body(default_factory=GenerateRequest),
store: DataProcessStore = Depends(get_data_process_store),
) -> dict[str, Any]:
return _start_generation(task_id, payload, background_tasks, store)
@router.post("/{task_id}/start")
def start(
task_id: str,
background_tasks: BackgroundTasks,
payload: GenerateRequest = Body(default_factory=GenerateRequest),
store: DataProcessStore = Depends(get_data_process_store),
storage: LocalDataProcessStorage = Depends(get_data_process_storage),
) -> dict[str, Any]:
with api_errors():
items = _prepare_preview_items(task_id, store, storage)
store.replace_preview_items(task_id, items)
return _start_generation(task_id, payload, background_tasks, store)
@router.post("/{task_id}/stop")
def stop(
task_id: str,
store: DataProcessStore = Depends(get_data_process_store),
) -> dict[str, Any]:
with api_errors():
store.stop_task(task_id)
return ok(store.progress(task_id), "data process task stopped")
@router.get("/{task_id}/progress")
def progress(
task_id: str,
store: DataProcessStore = Depends(get_data_process_store),
) -> dict[str, Any]:
with api_errors():
return ok(store.progress(task_id))
@router.get("/{task_id}/results")
def results(
task_id: str,
page: int = Query(default=1, ge=1),
page_size: int = Query(default=100, ge=1, le=1000),
status: Literal["valid", "modified", "invalid"] | None = Query(default=None),
split: Literal["train", "validation", "test"] | None = Query(default=None),
keyword: str | None = Query(default=None),
store: DataProcessStore = Depends(get_data_process_store),
) -> dict[str, Any]:
with api_errors():
return ok(
store.list_results(
task_id,
page=page,
page_size=page_size,
status=status,
split=split,
keyword=keyword,
)
)
@router.put("/{task_id}/results/{result_id}")
def update_result(
task_id: str,
result_id: str,
payload: ResultUpdate,
store: DataProcessStore = Depends(get_data_process_store),
) -> dict[str, Any]:
with api_errors():
task = store.get_task(task_id)
current = store.get_result(task_id, result_id)
update = payload.model_dump(exclude_unset=True, mode="json")
merged = {**current, **update}
preview_id = current.get("preview_item_id")
source_content = ""
if preview_id:
preview = store.get_preview_item(task_id, str(preview_id))
source_content = str(
preview.get("edited_content") or preview.get("original_content") or ""
)
minimum = max(
1,
int(
_value(
task.get("config") or {},
"min_output_length",
"minOutputLength",
20,
)
or 20
),
)
quality = score_quality(
merged,
min_output_length=minimum,
source_content=source_content,
)
update["quality_score"] = asdict(quality)
result = store.update_result(
task_id,
result_id,
update,
)
return ok(result, "data process result updated")
@router.post("/{task_id}/results/{result_id}/restore")
def restore_result(
task_id: str,
result_id: str,
store: DataProcessStore = Depends(get_data_process_store),
) -> dict[str, Any]:
with api_errors():
task = store.get_task(task_id)
current = store.get_result(task_id, result_id)
restored = {
**current,
"instruction": current.get("original_instruction") or current.get("instruction") or "",
"input": current.get("original_input") or current.get("input") or "",
"output": current.get("original_output") or current.get("output") or "",
}
preview_id = current.get("preview_item_id")
source_content = ""
if preview_id:
preview = store.get_preview_item(task_id, str(preview_id))
source_content = str(
preview.get("edited_content") or preview.get("original_content") or ""
)
minimum = max(
1,
int(
_value(
task.get("config") or {},
"min_output_length",
"minOutputLength",
20,
)
or 20
),
)
quality = score_quality(
restored,
min_output_length=minimum,
source_content=source_content,
)
restored = store.update_result(
task_id,
result_id,
{
"instruction": restored["instruction"],
"input": restored["input"],
"output": restored["output"],
"quality_score": asdict(quality),
"expected_updated_at": current.get("updated_at"),
},
)
return ok(restored, "data process result restored")
@router.post("/{task_id}/publish")
def publish(
task_id: str,
payload: PublishRequest,
store: DataProcessStore = Depends(get_data_process_store),
) -> dict[str, Any]:
with api_errors():
result = store.publish(task_id, payload.model_dump(mode="json"))
message = "dataset published" if result["created"] else "dataset already published"
return ok(result, message)