311 lines
11 KiB
Python
311 lines
11 KiB
Python
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import json
|
|||
|
|
import os
|
|||
|
|
from pathlib import Path
|
|||
|
|
from typing import Any
|
|||
|
|
|
|||
|
|
from fastapi import APIRouter, Body, UploadFile, File
|
|||
|
|
from fastapi.responses import FileResponse
|
|||
|
|
|
|||
|
|
from app.api.v1.endpoints.platform import ok, fail
|
|||
|
|
from app.db.platform_store import get_platform_store, new_id
|
|||
|
|
|
|||
|
|
|
|||
|
|
router = APIRouter(prefix="/data-convert", tags=["data-convert"])
|
|||
|
|
|
|||
|
|
# 存储根目录
|
|||
|
|
STORAGE_ROOT = Path(__file__).resolve().parents[3] / "storage" / "data-convert"
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _task_dir(task_id: str) -> Path:
|
|||
|
|
return STORAGE_ROOT / task_id
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _input_dir(task_id: str) -> Path:
|
|||
|
|
return _task_dir(task_id) / "input"
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _output_dir(task_id: str) -> Path:
|
|||
|
|
return _task_dir(task_id) / "output"
|
|||
|
|
|
|||
|
|
|
|||
|
|
@router.get("")
|
|||
|
|
def list_tasks(page: int = 1, page_size: int = 20) -> dict[str, Any]:
|
|||
|
|
store = get_platform_store()
|
|||
|
|
with store.connect() as conn:
|
|||
|
|
rows = conn.execute(
|
|||
|
|
"SELECT * FROM data_convert_tasks WHERE deleted_at IS NULL "
|
|||
|
|
"ORDER BY create_time DESC LIMIT %s OFFSET %s",
|
|||
|
|
(page_size, (page - 1) * page_size),
|
|||
|
|
).fetchall()
|
|||
|
|
total = conn.execute(
|
|||
|
|
"SELECT COUNT(*) FROM data_convert_tasks WHERE deleted_at IS NULL"
|
|||
|
|
).fetchone()[0]
|
|||
|
|
return ok({"items": [dict(r) for r in rows], "total": total})
|
|||
|
|
|
|||
|
|
|
|||
|
|
@router.post("")
|
|||
|
|
def create_task(payload: dict[str, Any] = Body(...)) -> dict[str, Any]:
|
|||
|
|
name = str(payload.get("name") or "").strip()
|
|||
|
|
if not name:
|
|||
|
|
raise fail(400, "name is required")
|
|||
|
|
task_id = new_id("dct")
|
|||
|
|
output_filename = str(payload.get("output_filename") or "converted-data.jsonl").strip()
|
|||
|
|
description = str(payload.get("description") or "").strip()
|
|||
|
|
store = get_platform_store()
|
|||
|
|
with store.connect() as conn:
|
|||
|
|
conn.execute(
|
|||
|
|
"INSERT INTO data_convert_tasks (id, name, description, output_filename) "
|
|||
|
|
"VALUES (%s, %s, %s, %s)",
|
|||
|
|
(task_id, name, description, output_filename),
|
|||
|
|
)
|
|||
|
|
# 创建目录
|
|||
|
|
_input_dir(task_id).mkdir(parents=True, exist_ok=True)
|
|||
|
|
_output_dir(task_id).mkdir(parents=True, exist_ok=True)
|
|||
|
|
return ok(_get_task(task_id))
|
|||
|
|
|
|||
|
|
|
|||
|
|
@router.get("/{task_id}")
|
|||
|
|
def get_task(task_id: str) -> dict[str, Any]:
|
|||
|
|
task = _get_task(task_id)
|
|||
|
|
if not task:
|
|||
|
|
raise fail(404, "task not found")
|
|||
|
|
# 附加输入文件列表
|
|||
|
|
input_dir = _input_dir(task_id)
|
|||
|
|
files = []
|
|||
|
|
if input_dir.exists():
|
|||
|
|
for f in sorted(input_dir.iterdir()):
|
|||
|
|
if f.is_file():
|
|||
|
|
files.append({"name": f.name, "size": f.stat().st_size})
|
|||
|
|
task["input_files"] = files
|
|||
|
|
return ok(task)
|
|||
|
|
|
|||
|
|
|
|||
|
|
@router.post("/{task_id}/source-files")
|
|||
|
|
async def upload_source_files(
|
|||
|
|
task_id: str,
|
|||
|
|
files: list[UploadFile] = File(...),
|
|||
|
|
) -> dict[str, Any]:
|
|||
|
|
task = _get_task(task_id)
|
|||
|
|
if not task:
|
|||
|
|
raise fail(404, "task not found")
|
|||
|
|
if task["status"] not in ("pending", "uploaded"):
|
|||
|
|
raise fail(400, "task is not editable")
|
|||
|
|
input_dir = _input_dir(task_id)
|
|||
|
|
input_dir.mkdir(parents=True, exist_ok=True)
|
|||
|
|
staged = []
|
|||
|
|
for upload in files:
|
|||
|
|
name = Path(upload.filename or "input.json").name
|
|||
|
|
if not name.lower().endswith(".json"):
|
|||
|
|
raise fail(415, f"only JSON files are supported: {name}")
|
|||
|
|
target = input_dir / name
|
|||
|
|
content = await upload.read()
|
|||
|
|
target.write_bytes(content)
|
|||
|
|
staged.append({"name": name, "size": len(content)})
|
|||
|
|
store = get_platform_store()
|
|||
|
|
# 标记上传完成
|
|||
|
|
with store.connect() as conn:
|
|||
|
|
conn.execute(
|
|||
|
|
"UPDATE data_convert_tasks SET status='uploaded', update_time=NOW() WHERE id=%s",
|
|||
|
|
(task_id,),
|
|||
|
|
)
|
|||
|
|
# 自动转换并导入数据集
|
|||
|
|
try:
|
|||
|
|
output_dir = _output_dir(task_id)
|
|||
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|||
|
|
output_path = output_dir / (task["output_filename"] or "converted-data.jsonl")
|
|||
|
|
# 清空旧输出（如果重新上传）
|
|||
|
|
if output_path.exists():
|
|||
|
|
output_path.unlink()
|
|||
|
|
input_count = 0
|
|||
|
|
output_count = 0
|
|||
|
|
for json_file in sorted(input_dir.iterdir()):
|
|||
|
|
if not json_file.is_file() or not json_file.name.lower().endswith(".json"):
|
|||
|
|
continue
|
|||
|
|
input_count += 1
|
|||
|
|
with open(json_file, "r", encoding="utf-8") as f:
|
|||
|
|
data = json.load(f)
|
|||
|
|
if isinstance(data, list):
|
|||
|
|
records = data
|
|||
|
|
elif isinstance(data, dict):
|
|||
|
|
records = [data]
|
|||
|
|
else:
|
|||
|
|
raise ValueError(f"JSON must be object or array: {json_file.name}")
|
|||
|
|
with open(output_path, "a", encoding="utf-8") as f:
|
|||
|
|
for record in records:
|
|||
|
|
f.write(json.dumps(record, ensure_ascii=False) + "\n")
|
|||
|
|
output_count += 1
|
|||
|
|
with store.connect() as conn:
|
|||
|
|
conn.execute(
|
|||
|
|
"UPDATE data_convert_tasks SET status='completed', "
|
|||
|
|
"input_count=%s, output_count=%s, update_time=NOW() WHERE id=%s",
|
|||
|
|
(input_count, output_count, task_id),
|
|||
|
|
)
|
|||
|
|
# 自动导入数据集
|
|||
|
|
content = output_path.read_text(encoding="utf-8")
|
|||
|
|
size_bytes = len(content.encode("utf-8"))
|
|||
|
|
dataset = store.create_dataset({
|
|||
|
|
"name": task["name"],
|
|||
|
|
"type": "train",
|
|||
|
|
"storage_type": "local",
|
|||
|
|
"source": "upload",
|
|||
|
|
"task_id": task_id,
|
|||
|
|
"size": f"{size_bytes} B",
|
|||
|
|
"count": output_count,
|
|||
|
|
"description": f"由数据类型转换任务 {task_id} 自动导入",
|
|||
|
|
})
|
|||
|
|
dataset_id = dataset["id"]
|
|||
|
|
with store.connect() as conn:
|
|||
|
|
store.add_dataset_file(conn, dataset_id, task["output_filename"] or "converted-data.jsonl", content)
|
|||
|
|
return ok({
|
|||
|
|
"staged_files": staged,
|
|||
|
|
"auto_converted": True,
|
|||
|
|
"dataset_id": dataset_id,
|
|||
|
|
"input_count": input_count,
|
|||
|
|
"output_count": output_count,
|
|||
|
|
})
|
|||
|
|
except Exception as exc:
|
|||
|
|
with store.connect() as conn:
|
|||
|
|
conn.execute(
|
|||
|
|
"UPDATE data_convert_tasks SET status='failed', error_message=%s, update_time=NOW() WHERE id=%s",
|
|||
|
|
(str(exc)[:500], task_id),
|
|||
|
|
)
|
|||
|
|
return ok({"staged_files": staged, "auto_converted": False, "error": str(exc)[:500]})
|
|||
|
|
|
|||
|
|
|
|||
|
|
@router.post("/{task_id}/run")
|
|||
|
|
def run_convert(task_id: str) -> dict[str, Any]:
|
|||
|
|
task = _get_task(task_id)
|
|||
|
|
if not task:
|
|||
|
|
raise fail(404, "task not found")
|
|||
|
|
if task["status"] not in ("uploaded", "completed", "failed"):
|
|||
|
|
raise fail(400, "please upload source files first")
|
|||
|
|
# 标记运行中
|
|||
|
|
store = get_platform_store()
|
|||
|
|
with store.connect() as conn:
|
|||
|
|
conn.execute(
|
|||
|
|
"UPDATE data_convert_tasks SET status='running', error_message='', update_time=NOW() WHERE id=%s",
|
|||
|
|
(task_id,),
|
|||
|
|
)
|
|||
|
|
try:
|
|||
|
|
input_dir = _input_dir(task_id)
|
|||
|
|
output_dir = _output_dir(task_id)
|
|||
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|||
|
|
output_path = output_dir / (task["output_filename"] or "converted-data.jsonl")
|
|||
|
|
input_count = 0
|
|||
|
|
output_count = 0
|
|||
|
|
for json_file in sorted(input_dir.iterdir()):
|
|||
|
|
if not json_file.is_file() or not json_file.name.lower().endswith(".json"):
|
|||
|
|
continue
|
|||
|
|
input_count += 1
|
|||
|
|
with open(json_file, "r", encoding="utf-8") as f:
|
|||
|
|
data = json.load(f)
|
|||
|
|
if isinstance(data, list):
|
|||
|
|
records = data
|
|||
|
|
elif isinstance(data, dict):
|
|||
|
|
records = [data]
|
|||
|
|
else:
|
|||
|
|
raise ValueError(f"JSON must be object or array: {json_file.name}")
|
|||
|
|
with open(output_path, "a", encoding="utf-8") as f:
|
|||
|
|
for record in records:
|
|||
|
|
f.write(json.dumps(record, ensure_ascii=False) + "\n")
|
|||
|
|
output_count += 1
|
|||
|
|
# 更新任务状态
|
|||
|
|
with store.connect() as conn:
|
|||
|
|
conn.execute(
|
|||
|
|
"UPDATE data_convert_tasks SET status='completed', "
|
|||
|
|
"input_count=%s, output_count=%s, update_time=NOW() WHERE id=%s",
|
|||
|
|
(input_count, output_count, task_id),
|
|||
|
|
)
|
|||
|
|
except Exception as exc:
|
|||
|
|
with store.connect() as conn:
|
|||
|
|
conn.execute(
|
|||
|
|
"UPDATE data_convert_tasks SET status='failed', error_message=%s, update_time=NOW() WHERE id=%s",
|
|||
|
|
(str(exc)[:500], task_id),
|
|||
|
|
)
|
|||
|
|
raise fail(500, f"convert failed: {exc}")
|
|||
|
|
return ok(_get_task(task_id))
|
|||
|
|
|
|||
|
|
|
|||
|
|
@router.get("/{task_id}/download")
|
|||
|
|
def download_result(task_id: str):
|
|||
|
|
task = _get_task(task_id)
|
|||
|
|
if not task:
|
|||
|
|
raise fail(404, "task not found")
|
|||
|
|
if task["status"] != "completed":
|
|||
|
|
raise fail(400, "task is not completed")
|
|||
|
|
output_path = _output_dir(task_id) / (task["output_filename"] or "converted-data.jsonl")
|
|||
|
|
if not output_path.exists():
|
|||
|
|
raise fail(404, "output file not found")
|
|||
|
|
return FileResponse(
|
|||
|
|
str(output_path),
|
|||
|
|
media_type="application/octet-stream",
|
|||
|
|
filename=task["output_filename"] or "converted-data.jsonl",
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
@router.post("/{task_id}/import-as-dataset")
|
|||
|
|
def import_as_dataset(
|
|||
|
|
task_id: str,
|
|||
|
|
payload: dict[str, Any] = Body(default={}),
|
|||
|
|
) -> dict[str, Any]:
|
|||
|
|
"""把已转换的 JSONL 文件导入为数据集管理中的上传任务记录（source='task'）。"""
|
|||
|
|
task = _get_task(task_id)
|
|||
|
|
if not task:
|
|||
|
|
raise fail(404, "task not found")
|
|||
|
|
if task["status"] != "completed":
|
|||
|
|
raise fail(400, "task is not completed")
|
|||
|
|
output_path = _output_dir(task_id) / (task["output_filename"] or "converted-data.jsonl")
|
|||
|
|
if not output_path.exists():
|
|||
|
|
raise fail(404, "output file not found")
|
|||
|
|
content = output_path.read_text(encoding="utf-8")
|
|||
|
|
dataset_name = str(payload.get("name") or task["name"]).strip()
|
|||
|
|
description = str(payload.get("description") or f"由数据类型转换任务 {task_id} 导入").strip()
|
|||
|
|
size_bytes = len(content.encode("utf-8"))
|
|||
|
|
store = get_platform_store()
|
|||
|
|
# 用 store 提供的接口创建数据集与文件
|
|||
|
|
dataset = store.create_dataset({
|
|||
|
|
"name": dataset_name,
|
|||
|
|
"type": "train",
|
|||
|
|
"storage_type": "local",
|
|||
|
|
"source": "upload",
|
|||
|
|
"task_id": task_id,
|
|||
|
|
"size": f"{size_bytes} B",
|
|||
|
|
"count": task["output_count"],
|
|||
|
|
"description": description,
|
|||
|
|
})
|
|||
|
|
dataset_id = dataset["id"]
|
|||
|
|
with store.connect() as conn:
|
|||
|
|
store.add_dataset_file(conn, dataset_id, task["output_filename"] or "converted-data.jsonl", content)
|
|||
|
|
return ok({"dataset_id": dataset_id, "name": dataset_name})
|
|||
|
|
|
|||
|
|
|
|||
|
|
@router.delete("/{task_id}")
|
|||
|
|
def delete_task(task_id: str) -> dict[str, Any]:
|
|||
|
|
task = _get_task(task_id)
|
|||
|
|
if not task:
|
|||
|
|
raise fail(404, "task not found")
|
|||
|
|
store = get_platform_store()
|
|||
|
|
with store.connect() as conn:
|
|||
|
|
conn.execute(
|
|||
|
|
"UPDATE data_convert_tasks SET deleted_at=NOW() WHERE id=%s",
|
|||
|
|
(task_id,),
|
|||
|
|
)
|
|||
|
|
# 清理文件
|
|||
|
|
import shutil
|
|||
|
|
task_dir = _task_dir(task_id)
|
|||
|
|
if task_dir.exists():
|
|||
|
|
shutil.rmtree(task_dir, ignore_errors=True)
|
|||
|
|
return ok({"deleted": task_id})
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _get_task(task_id: str) -> dict[str, Any] | None:
|
|||
|
|
store = get_platform_store()
|
|||
|
|
with store.connect() as conn:
|
|||
|
|
row = conn.execute(
|
|||
|
|
"SELECT * FROM data_convert_tasks WHERE id=%s AND deleted_at IS NULL",
|
|||
|
|
(task_id,),
|
|||
|
|
).fetchone()
|
|||
|
|
return dict(row) if row else None
|