1570 lines
57 KiB
Python
1570 lines
57 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)
|
|
)
|
|
if generation_model:
|
|
runtime_config = {
|
|
**config,
|
|
"generation_prompt": _value(
|
|
config, "generation_prompt", "generationPrompt", ""
|
|
),
|
|
"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,
|
|
)
|
|
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():
|
|
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)
|