1565 lines
56 KiB
Python
1565 lines
56 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,
|
||
TextChunk,
|
||
canonical_record_json,
|
||
chunk_unstructured,
|
||
content_quality_flags,
|
||
desensitize_pii,
|
||
desensitize_structured_record,
|
||
detect_document_structure,
|
||
estimate_token_count,
|
||
extract_pdf_page_texts,
|
||
generate_standard_records,
|
||
is_near_duplicate,
|
||
near_duplicate_fingerprint,
|
||
parse_text_content,
|
||
preprocess_structured_records,
|
||
score_quality,
|
||
)
|
||
from app.modules.data_process.generation import generate_model_records
|
||
from app.modules.data_process.storage import (
|
||
LocalDataProcessStorage,
|
||
StagedSourceObject,
|
||
get_data_process_storage,
|
||
)
|
||
from app.modules.data_process.store import (
|
||
ConflictError,
|
||
DataProcessStore,
|
||
DataProcessStoreError,
|
||
InvalidStateError,
|
||
NotFoundError,
|
||
get_data_process_store,
|
||
new_id,
|
||
)
|
||
from app.schemas.data_process import (
|
||
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 _shift_chunk(chunk: TextChunk, offset: int, source_text: str) -> TextChunk:
|
||
start = offset + chunk.start
|
||
end = offset + chunk.end
|
||
return TextChunk(
|
||
content=chunk.content,
|
||
start=start,
|
||
end=end,
|
||
start_line=source_text.count("\n", 0, start) + 1,
|
||
end_line=source_text.count("\n", 0, max(start, end - 1)) + 1,
|
||
token_count=chunk.token_count,
|
||
)
|
||
|
||
|
||
def _merge_short_chunks(
|
||
chunks: list[TextChunk],
|
||
source_text: str,
|
||
*,
|
||
min_token_count: int,
|
||
chunk_size: int,
|
||
) -> list[TextChunk]:
|
||
"""按原文顺序合并相邻短块，同时严格遵守切片大小上限。"""
|
||
|
||
merged: list[TextChunk] = []
|
||
index = 0
|
||
while index < len(chunks):
|
||
current = chunks[index]
|
||
if current.token_count >= min_token_count:
|
||
merged.append(current)
|
||
index += 1
|
||
continue
|
||
|
||
candidates: list[tuple[int, int, int]] = []
|
||
if merged:
|
||
candidates.append((merged[-1].start, current.end, -1))
|
||
if index + 1 < len(chunks):
|
||
candidates.append((current.start, chunks[index + 1].end, 1))
|
||
selected = next(
|
||
(
|
||
(start, end, direction)
|
||
for start, end, direction in candidates
|
||
if estimate_token_count(source_text[start:end]) <= chunk_size
|
||
),
|
||
None,
|
||
)
|
||
if selected is None:
|
||
merged.append(current)
|
||
index += 1
|
||
continue
|
||
|
||
start, end, direction = selected
|
||
combined = TextChunk(
|
||
content=source_text[start:end],
|
||
start=start,
|
||
end=end,
|
||
start_line=source_text.count("\n", 0, start) + 1,
|
||
end_line=source_text.count("\n", 0, max(start, end - 1)) + 1,
|
||
token_count=estimate_token_count(source_text[start:end]),
|
||
)
|
||
if direction < 0:
|
||
merged[-1] = combined
|
||
index += 1
|
||
else:
|
||
merged.append(combined)
|
||
index += 2
|
||
return merged
|
||
|
||
|
||
def _chunk_source_text(
|
||
text: str,
|
||
config: dict[str, Any],
|
||
preprocess_options: set[str],
|
||
) -> list[tuple[TextChunk, tuple[str, ...]]]:
|
||
"""按可选文档结构分段后切片，结构边界之间不共享 overlap。"""
|
||
|
||
method = str(_value(config, "chunk_method", "chunkMethod", "structure"))
|
||
detect_structure = (
|
||
method == "structure" or "detect_document_structure" in preprocess_options
|
||
)
|
||
preserve_context = "preserve_context" in preprocess_options
|
||
merge_short = "merge_short_content" in preprocess_options
|
||
chunk_size = int(_value(config, "chunk_size", "chunkSize", 800))
|
||
configured_minimum = int(_value(config, "min_chunk_size", "minChunkSize", 100))
|
||
minimum = configured_minimum if merge_short else 1
|
||
overlap = (
|
||
int(_value(config, "chunk_overlap", "chunkOverlap", 100))
|
||
if preserve_context
|
||
else 0
|
||
)
|
||
common = {
|
||
"method": method,
|
||
"chunk_size": chunk_size,
|
||
"chunk_overlap": overlap,
|
||
"min_chunk_size": minimum,
|
||
"custom_delimiter": str(
|
||
_value(config, "custom_delimiter", "customDelimiter", "") or ""
|
||
),
|
||
"preserve_code_blocks": bool(
|
||
_value(config, "preserve_code_blocks", "preserveCodeBlocks", False)
|
||
),
|
||
"preserve_tables": bool(
|
||
_value(config, "preserve_tables", "preserveTables", False)
|
||
),
|
||
"preserve_lists": bool(
|
||
_value(config, "preserve_lists", "preserveLists", False)
|
||
),
|
||
}
|
||
|
||
sections: list[tuple[int, int, tuple[str, ...]]] = [(0, len(text), ())]
|
||
if detect_structure:
|
||
structure = detect_document_structure(text)
|
||
if structure.headings:
|
||
sections = []
|
||
first_start = structure.headings[0].start
|
||
if first_start > 0 and text[:first_start].strip():
|
||
sections.append((0, first_start, ()))
|
||
stack: list[tuple[int, str]] = []
|
||
for index, heading in enumerate(structure.headings):
|
||
while stack and stack[-1][0] >= heading.level:
|
||
stack.pop()
|
||
stack.append((heading.level, heading.title))
|
||
end = (
|
||
structure.headings[index + 1].start
|
||
if index + 1 < len(structure.headings)
|
||
else len(text)
|
||
)
|
||
sections.append((heading.start, end, tuple(title for _, title in stack)))
|
||
|
||
result: list[tuple[TextChunk, tuple[str, ...]]] = []
|
||
for start, end, heading_path in sections:
|
||
section_text = text[start:end]
|
||
local_chunks = chunk_unstructured(section_text, **common)
|
||
shifted = [_shift_chunk(chunk, start, text) for chunk in local_chunks]
|
||
result.extend((chunk, heading_path) for chunk in shifted)
|
||
|
||
if merge_short and result:
|
||
# 结构分段只负责提供标题路径和隔离 overlap，不应让目录项或短小节
|
||
# 突破 min_chunk_size 约束。合并后保留首个原始块的标题路径。
|
||
heading_paths = {chunk.start: heading_path for chunk, heading_path in result}
|
||
merged = _merge_short_chunks(
|
||
[chunk for chunk, _ in result],
|
||
text,
|
||
min_token_count=configured_minimum,
|
||
chunk_size=chunk_size,
|
||
)
|
||
result = [(chunk, heading_paths.get(chunk.start, ())) for chunk in merged]
|
||
return result
|
||
|
||
|
||
_NEGATION_MARKERS = frozenset({"不", "无", "未", "否", "没有", "并非", "not", "no", "never"})
|
||
|
||
|
||
def _safe_near_duplicate(left: str, right: str) -> bool:
|
||
"""保守判断近重复，数字或否定含义变化时始终保留。"""
|
||
|
||
if min(estimate_token_count(left), estimate_token_count(right)) < 20:
|
||
return False
|
||
if re.findall(r"\d+(?:\.\d+)?", left) != re.findall(r"\d+(?:\.\d+)?", right):
|
||
return False
|
||
left_lower = left.casefold()
|
||
right_lower = right.casefold()
|
||
left_negations = {marker for marker in _NEGATION_MARKERS if marker in left_lower}
|
||
right_negations = {marker for marker in _NEGATION_MARKERS if marker in right_lower}
|
||
if left_negations != right_negations:
|
||
return False
|
||
return is_near_duplicate(
|
||
left,
|
||
right,
|
||
similarity_threshold=0.92,
|
||
max_hamming_distance=2,
|
||
)
|
||
|
||
|
||
def _near_duplicate_band_keys(content: str) -> tuple[tuple[int, int], ...]:
|
||
"""将 64 位 SimHash 分为三段，汉明距离不超过 2 时至少命中一段。"""
|
||
|
||
fingerprint = int(near_duplicate_fingerprint(content), 16)
|
||
widths = (22, 21, 21)
|
||
shift = 0
|
||
keys: list[tuple[int, int]] = []
|
||
for index, width in enumerate(widths):
|
||
keys.append((index, (fingerprint >> shift) & ((1 << width) - 1)))
|
||
shift += width
|
||
return tuple(keys)
|
||
|
||
|
||
def _build_preview_items(
|
||
task: dict[str, Any], source_files: list[dict[str, Any]]
|
||
) -> list[dict[str, Any]]:
|
||
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":
|
||
chunks = _chunk_source_text(parsed.text, config, preprocess_options)
|
||
for chunk, heading_path in chunks:
|
||
preprocess_flags = content_quality_flags(
|
||
chunk.content,
|
||
min_chars=0,
|
||
min_tokens=0,
|
||
)
|
||
flag_set = set(preprocess_flags)
|
||
if "clean_invalid_content" in preprocess_options and flag_set & {
|
||
"empty_content",
|
||
"low_printable_ratio",
|
||
"repetitive_content",
|
||
}:
|
||
continue
|
||
if "filter_low_quality" in preprocess_options and flag_set & {
|
||
"content_too_long",
|
||
"mojibake",
|
||
"low_printable_ratio",
|
||
"repetitive_content",
|
||
}:
|
||
continue
|
||
if "deduplicate_content" in preprocess_options:
|
||
band_keys = _near_duplicate_band_keys(chunk.content)
|
||
candidates = {
|
||
previous
|
||
for key in band_keys
|
||
for previous in seen_near_duplicate_bands.get(key, ())
|
||
}
|
||
if any(
|
||
_safe_near_duplicate(chunk.content, previous)
|
||
for previous in candidates
|
||
):
|
||
continue
|
||
for key in band_keys:
|
||
seen_near_duplicate_bands.setdefault(key, []).append(chunk.content)
|
||
content = chunk.content
|
||
pii_counts: dict[str, int] = {}
|
||
if should_desensitize:
|
||
content, pii_counts = desensitize_pii(content)
|
||
quality = _preview_quality(content, config)
|
||
quality["pii_replacements"] = pii_counts
|
||
quality["preprocess_flags"] = list(preprocess_flags)
|
||
chunk_method = str(
|
||
_value(config, "chunk_method", "chunkMethod", "structure")
|
||
)
|
||
if (
|
||
chunk_method == "structure"
|
||
or "detect_document_structure" in preprocess_options
|
||
):
|
||
quality["heading_path"] = list(heading_path)
|
||
append_item(
|
||
{
|
||
"source_file_id": source["id"],
|
||
"original_content": chunk.content,
|
||
"edited_content": content,
|
||
"source_start": chunk.start,
|
||
"source_end": chunk.end,
|
||
"source_start_line": chunk.start_line,
|
||
"source_end_line": chunk.end_line,
|
||
"token_count": chunk.token_count,
|
||
"status": "modified" if content != chunk.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)
|
||
task["source_files"] = store.list_source_files(task_id)
|
||
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.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,
|
||
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")
|
||
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),
|
||
) -> dict[str, Any]:
|
||
with api_errors():
|
||
selected_ids = payload.source_file_ids
|
||
items = _prepare_preview_items(task_id, store, selected_ids)
|
||
created = store.replace_preview_items(
|
||
task_id,
|
||
items,
|
||
source_file_ids=selected_ids,
|
||
)
|
||
target_ids = selected_ids or [
|
||
str(source["id"]) for source in store.list_source_files(task_id)
|
||
]
|
||
file_counts = dict.fromkeys(target_ids, 0)
|
||
for item in created:
|
||
source_file_id = str(item.get("source_file_id") or "")
|
||
if source_file_id in file_counts:
|
||
file_counts[source_file_id] += 1
|
||
files = [
|
||
{
|
||
"source_file_id": source_file_id,
|
||
"preview_count": count,
|
||
"status": "completed" if count else "empty",
|
||
}
|
||
for source_file_id, count in file_counts.items()
|
||
]
|
||
return ok(
|
||
{
|
||
"items": created,
|
||
"total": len(created),
|
||
"page": 1,
|
||
"page_size": len(created),
|
||
"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),
|
||
) -> dict[str, Any]:
|
||
with api_errors():
|
||
items = _prepare_preview_items(task_id, store)
|
||
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)
|