2093 lines
87 KiB
Python
2093 lines
87 KiB
Python
from __future__ import annotations
|
||
|
||
import hashlib
|
||
import json
|
||
import re
|
||
import uuid
|
||
from collections.abc import Iterator, Sequence
|
||
from contextlib import contextmanager
|
||
from datetime import UTC, date, datetime
|
||
from decimal import Decimal
|
||
from functools import lru_cache
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import psycopg
|
||
from psycopg.rows import dict_row
|
||
|
||
from app.core.config import get_settings
|
||
from app.modules.data_process.algorithms import estimate_token_count, stable_split_assignments
|
||
|
||
TASK_STATUSES = {"pending", "running", "completed", "failed", "stopped"}
|
||
EDITABLE_STATUSES = {"pending", "failed", "stopped", "completed"}
|
||
|
||
_PREVIEW_CONFIG_ALIASES = {
|
||
"preprocess_options": "preprocessOptions",
|
||
"chunk_method": "chunkMethod",
|
||
"chunk_size": "chunkSize",
|
||
"chunk_overlap": "chunkOverlap",
|
||
"min_chunk_size": "minChunkSize",
|
||
"semantic_breakpoint_percentile": "semanticBreakpointPercentile",
|
||
"preserve_tables": "preserveTables",
|
||
"preserve_code_blocks": "preserveCodeBlocks",
|
||
"preserve_lists": "preserveLists",
|
||
}
|
||
_UNSTRUCTURED_PREVIEW_DEFAULTS: dict[str, Any] = {
|
||
"chunk_method": "layout_hybrid",
|
||
"chunk_size": 800,
|
||
"chunk_overlap": 100,
|
||
"min_chunk_size": 100,
|
||
"semantic_breakpoint_percentile": 95,
|
||
"preserve_tables": True,
|
||
"preserve_code_blocks": True,
|
||
"preserve_lists": True,
|
||
}
|
||
_REGENERATION_MARKER_KEY = "_regeneration_prepared"
|
||
|
||
|
||
class DataProcessStoreError(RuntimeError):
|
||
pass
|
||
|
||
|
||
class NotFoundError(DataProcessStoreError):
|
||
pass
|
||
|
||
|
||
class ConflictError(DataProcessStoreError):
|
||
pass
|
||
|
||
|
||
class InvalidStateError(DataProcessStoreError):
|
||
pass
|
||
|
||
|
||
def utcnow() -> str:
|
||
return datetime.now(UTC).replace(microsecond=0).isoformat().replace("+00:00", "Z")
|
||
|
||
|
||
def new_id(prefix: str) -> str:
|
||
return f"{prefix}_{uuid.uuid4().hex[:20]}"
|
||
|
||
|
||
def json_dumps(value: Any) -> str:
|
||
return json.dumps(value, ensure_ascii=False, separators=(",", ":"))
|
||
|
||
|
||
def _database_url(value: str) -> str:
|
||
return value.replace("postgresql+psycopg://", "postgresql://")
|
||
|
||
|
||
def _json_value(value: Any, default: Any) -> Any:
|
||
if value is None or value == "":
|
||
return default
|
||
if isinstance(value, (dict, list)):
|
||
return value
|
||
try:
|
||
return json.loads(value)
|
||
except (TypeError, json.JSONDecodeError):
|
||
return default
|
||
|
||
|
||
def _task_output_type(task: dict[str, Any]) -> str:
|
||
config = _json_value(task.get("config"), {})
|
||
if not isinstance(config, dict):
|
||
return "standard"
|
||
return str(config.get("output_type") or config.get("outputType") or "standard")
|
||
|
||
|
||
def _task_reasoning_detail(task: dict[str, Any]) -> str:
|
||
config = _json_value(task.get("config"), {})
|
||
if not isinstance(config, dict):
|
||
return "normal"
|
||
return str(
|
||
config.get("reasoning_detail")
|
||
or config.get("reasoningDetail")
|
||
or "normal"
|
||
)
|
||
|
||
|
||
def _reasoning_output_is_valid(value: Any) -> bool:
|
||
match = re.fullmatch(
|
||
r"\s*<think>\s*(?P<reasoning>[\s\S]*?)\s*</think>\s*(?P<answer>[\s\S]+?)\s*",
|
||
str(value or ""),
|
||
flags=re.IGNORECASE,
|
||
)
|
||
return bool(
|
||
match
|
||
and match.group("reasoning").strip()
|
||
and match.group("answer").strip()
|
||
and all(
|
||
tag not in part.lower()
|
||
for tag in ("<think", "</think")
|
||
for part in (match.group("reasoning"), match.group("answer"))
|
||
)
|
||
)
|
||
|
||
|
||
def _preview_config_value(config: dict[str, Any], key: str, default: Any) -> Any:
|
||
if key in config:
|
||
return config[key]
|
||
return config.get(_PREVIEW_CONFIG_ALIASES[key], default)
|
||
|
||
|
||
def _normalized_preprocess_options(config: dict[str, Any]) -> Any:
|
||
value = _preview_config_value(config, "preprocess_options", [])
|
||
if isinstance(value, (list, tuple, set)):
|
||
return tuple(sorted({str(item) for item in value}))
|
||
return value
|
||
|
||
|
||
def _preview_config_projection(
|
||
process_type: str,
|
||
config: dict[str, Any],
|
||
) -> dict[str, Any]:
|
||
"""只投影会改变预览切片的配置。
|
||
|
||
生成模型、提示词、温度等参数不影响源文切片，因此不应该
|
||
破坏用户已经校对过的预览内容。
|
||
"""
|
||
|
||
projection: dict[str, Any] = {
|
||
"preprocess_options": _normalized_preprocess_options(config),
|
||
}
|
||
if process_type != "unstructured":
|
||
return projection
|
||
for key, default in _UNSTRUCTURED_PREVIEW_DEFAULTS.items():
|
||
projection[key] = _preview_config_value(config, key, default)
|
||
return projection
|
||
|
||
|
||
def _preview_config_changed(
|
||
process_type: str,
|
||
current_config: dict[str, Any],
|
||
next_config: dict[str, Any],
|
||
) -> bool:
|
||
return _preview_config_projection(process_type, current_config) != _preview_config_projection(
|
||
process_type, next_config
|
||
)
|
||
|
||
|
||
def _regeneration_marker(task: dict[str, Any]) -> dict[str, Any] | None:
|
||
config = task.get("config")
|
||
if not isinstance(config, dict):
|
||
return None
|
||
marker = config.get(_REGENERATION_MARKER_KEY)
|
||
if not isinstance(marker, dict) or marker.get("prepared") is not True:
|
||
return None
|
||
return marker
|
||
|
||
|
||
def _is_regeneration_prepared(task: dict[str, Any]) -> bool:
|
||
return _regeneration_marker(task) is not None
|
||
|
||
|
||
def _public_task(item: dict[str, Any] | None) -> dict[str, Any] | None:
|
||
"""从 API 任务快照中移除服务端内部重新生成标记。"""
|
||
|
||
if item is None:
|
||
return None
|
||
public = dict(item)
|
||
config = public.get("config")
|
||
if isinstance(config, dict) and _REGENERATION_MARKER_KEY in config:
|
||
public["config"] = {
|
||
key: value for key, value in config.items() if key != _REGENERATION_MARKER_KEY
|
||
}
|
||
return public
|
||
|
||
|
||
def _serialize_value(value: Any) -> Any:
|
||
if isinstance(value, (datetime, date)):
|
||
return value.isoformat().replace("+00:00", "Z")
|
||
if isinstance(value, Decimal):
|
||
return float(value)
|
||
return value
|
||
|
||
|
||
def _source_storage_descriptor(
|
||
payload: dict[str, Any],
|
||
task_id: str,
|
||
file_id: str,
|
||
) -> tuple[str, dict[str, Any]]:
|
||
storage_object_id = str(
|
||
payload.get("storage_object_id")
|
||
or f"db://data-process/{task_id}/{file_id}/v1"
|
||
)
|
||
expected_local_prefix = f"local://data-process/{task_id}/{file_id}/v1/"
|
||
expected_database_reference = f"db://data-process/{task_id}/{file_id}/v1"
|
||
if storage_object_id.startswith(expected_local_prefix) and len(storage_object_id) > len(
|
||
expected_local_prefix
|
||
):
|
||
storage_backend = "local"
|
||
elif storage_object_id == expected_database_reference:
|
||
storage_backend = "database"
|
||
elif storage_object_id.startswith(("local://data-process/", "db://data-process/")):
|
||
raise DataProcessStoreError("source storage object owner mismatch")
|
||
else:
|
||
raise DataProcessStoreError("unsupported source storage object reference")
|
||
metadata = {
|
||
**(payload.get("metadata") or {}),
|
||
"storage_backend": storage_backend,
|
||
}
|
||
return storage_object_id, metadata
|
||
|
||
|
||
def _decode_row(row: dict[str, Any] | None) -> dict[str, Any] | None:
|
||
if row is None:
|
||
return None
|
||
item = {key: _serialize_value(value) for key, value in row.items()}
|
||
for key, default in {
|
||
"config": {},
|
||
"metadata": {},
|
||
"quality_score": {},
|
||
"versions": [],
|
||
"output_datasets": [],
|
||
}.items():
|
||
if key in item:
|
||
item[key] = _json_value(item[key], default)
|
||
return item
|
||
|
||
|
||
class DataProcessStore:
|
||
"""数据处理持久层。
|
||
|
||
构造函数不会连接数据库或执行迁移。部署方必须显式执行 002 SQL，
|
||
或在受控的管理命令中调用 :meth:`ensure_schema`，避免应用启动时
|
||
修改远程数据库。
|
||
"""
|
||
|
||
def __init__(self, database_url: str | None = None) -> None:
|
||
self.database_url = _database_url(database_url or get_settings().database_url)
|
||
|
||
@contextmanager
|
||
def connect(self) -> Iterator[psycopg.Connection[dict[str, Any]]]:
|
||
with psycopg.connect(self.database_url, row_factory=dict_row) as conn:
|
||
try:
|
||
yield conn
|
||
conn.commit()
|
||
except Exception:
|
||
conn.rollback()
|
||
raise
|
||
|
||
def ensure_schema(self) -> None:
|
||
"""显式安装数据处理表；API 路由和应用启动流程不会调用此方法。"""
|
||
schema_path = Path(__file__).resolve().parents[2] / "db" / "sql" / "002_data_process.sql"
|
||
sql = schema_path.read_text(encoding="utf-8")
|
||
with self.connect() as conn, conn.cursor() as cursor:
|
||
cursor.execute(sql)
|
||
|
||
def list_tasks(
|
||
self,
|
||
*,
|
||
page: int = 1,
|
||
page_size: int = 20,
|
||
keyword: str | None = None,
|
||
status: str | None = None,
|
||
process_type: str | None = None,
|
||
tenant_id: str | None = None,
|
||
project_id: str | None = None,
|
||
) -> dict[str, Any]:
|
||
clauses = ["task.deleted_at IS NULL"]
|
||
params: list[Any] = []
|
||
if keyword:
|
||
clauses.append("(task.name ILIKE %s OR COALESCE(task.description, '') ILIKE %s)")
|
||
pattern = f"%{keyword.strip()}%"
|
||
params.extend([pattern, pattern])
|
||
if status:
|
||
clauses.append("task.status = %s")
|
||
params.append(status)
|
||
if process_type:
|
||
clauses.append("task.process_type = %s")
|
||
params.append(process_type)
|
||
if tenant_id:
|
||
clauses.append("task.tenant_id = %s")
|
||
params.append(tenant_id)
|
||
if project_id:
|
||
clauses.append("task.project_id = %s")
|
||
params.append(project_id)
|
||
where = " AND ".join(clauses)
|
||
with self.connect() as conn:
|
||
total = conn.execute(
|
||
f"SELECT COUNT(*) AS count FROM data_process_tasks task WHERE {where}",
|
||
params,
|
||
).fetchone()["count"]
|
||
rows = conn.execute(
|
||
f"""
|
||
SELECT task.*,
|
||
(SELECT COUNT(*) FROM data_process_source_files source_file
|
||
WHERE source_file.task_id=task.id
|
||
AND source_file.deleted_at IS NULL) AS source_file_count
|
||
FROM data_process_tasks task
|
||
WHERE {where}
|
||
ORDER BY task.created_at DESC, task.id DESC
|
||
LIMIT %s OFFSET %s
|
||
""",
|
||
[*params, page_size, (page - 1) * page_size],
|
||
).fetchall()
|
||
return {
|
||
"items": [_public_task(_decode_row(row)) for row in rows],
|
||
"total": int(total),
|
||
"page": page,
|
||
"page_size": page_size,
|
||
}
|
||
|
||
def create_task(self, payload: dict[str, Any]) -> dict[str, Any]:
|
||
task_id = new_id("dpt")
|
||
now = utcnow()
|
||
try:
|
||
with self.connect() as conn:
|
||
row = conn.execute(
|
||
"""
|
||
INSERT INTO data_process_tasks
|
||
(id, name, description, status, process_type, source_dataset_id, config,
|
||
progress, tenant_id, project_id, owner_id, created_by, updated_by,
|
||
created_at, updated_at)
|
||
VALUES (%s, %s, %s, 'pending', %s, %s, %s, 0, %s, %s, %s, %s, %s, %s, %s)
|
||
RETURNING *
|
||
""",
|
||
(
|
||
task_id,
|
||
payload["name"],
|
||
payload.get("description") or "",
|
||
payload["process_type"],
|
||
payload.get("source_dataset_id"),
|
||
json_dumps(
|
||
{
|
||
key: value
|
||
for key, value in (payload.get("config") or {}).items()
|
||
if key != _REGENERATION_MARKER_KEY
|
||
}
|
||
),
|
||
payload.get("tenant_id"),
|
||
payload.get("project_id"),
|
||
payload.get("owner_id"),
|
||
payload.get("created_by"),
|
||
payload.get("created_by"),
|
||
now,
|
||
now,
|
||
),
|
||
).fetchone()
|
||
except psycopg.errors.UniqueViolation as exc:
|
||
raise ConflictError("data process task name already exists") from exc
|
||
return _public_task(_decode_row(row)) or {}
|
||
|
||
def get_task(self, task_id: str, *, for_update: bool = False) -> dict[str, Any]:
|
||
lock = " FOR UPDATE" if for_update else ""
|
||
with self.connect() as conn:
|
||
if for_update:
|
||
row = conn.execute(
|
||
f"SELECT * FROM data_process_tasks WHERE id=%s AND deleted_at IS NULL{lock}",
|
||
(task_id,),
|
||
).fetchone()
|
||
else:
|
||
row = conn.execute(
|
||
"""
|
||
SELECT task.*,
|
||
(SELECT COUNT(*) FROM data_process_source_files source
|
||
WHERE source.task_id=task.id AND source.deleted_at IS NULL)
|
||
AS source_file_count,
|
||
(SELECT COUNT(*) FROM data_process_preview_items preview
|
||
WHERE preview.task_id=task.id) AS preview_count,
|
||
(SELECT COALESCE(json_agg(json_build_object(
|
||
'id', dataset.id,
|
||
'name', dataset.name,
|
||
'type', dataset.type,
|
||
'count', dataset.count,
|
||
'dataset_split', CASE dataset.type
|
||
WHEN 'train' THEN 'train'
|
||
WHEN 'val' THEN 'validation'
|
||
WHEN 'test' THEN 'test'
|
||
ELSE NULL
|
||
END
|
||
) ORDER BY CASE dataset.type
|
||
WHEN 'train' THEN 1 WHEN 'val' THEN 2 WHEN 'test' THEN 3 ELSE 4 END), '[]'::json)
|
||
FROM datasets dataset
|
||
WHERE dataset.source='task'
|
||
AND dataset.deleted_at IS NULL
|
||
AND (
|
||
dataset.source_task_id=task.id
|
||
OR (dataset.source_task_id IS NULL AND dataset.task_id=task.id)
|
||
))
|
||
AS output_datasets,
|
||
CASE
|
||
WHEN task.started_at IS NOT NULL AND task.completed_at IS NOT NULL
|
||
THEN EXTRACT(EPOCH FROM (task.completed_at - task.started_at))
|
||
ELSE NULL
|
||
END AS duration_seconds
|
||
FROM data_process_tasks task
|
||
WHERE task.id=%s AND task.deleted_at IS NULL
|
||
""",
|
||
(task_id,),
|
||
).fetchone()
|
||
if not row:
|
||
raise NotFoundError("data process task not found")
|
||
return _public_task(_decode_row(row)) or {}
|
||
|
||
def _task_in_connection(
|
||
self,
|
||
conn: psycopg.Connection[dict[str, Any]],
|
||
task_id: str,
|
||
*,
|
||
for_update: bool = False,
|
||
) -> dict[str, Any]:
|
||
lock = " FOR UPDATE" if for_update else ""
|
||
row = conn.execute(
|
||
f"SELECT * FROM data_process_tasks WHERE id=%s AND deleted_at IS NULL{lock}",
|
||
(task_id,),
|
||
).fetchone()
|
||
if not row:
|
||
raise NotFoundError("data process task not found")
|
||
return _decode_row(row) or {}
|
||
|
||
@staticmethod
|
||
def _ensure_editable(task: dict[str, Any]) -> None:
|
||
if task["status"] not in EDITABLE_STATUSES:
|
||
raise InvalidStateError(f"task cannot be edited while status is {task['status']}")
|
||
if task.get("output_dataset_id") and not _is_regeneration_prepared(task):
|
||
raise InvalidStateError("published task cannot be edited")
|
||
|
||
def recover_legacy_aborted_regeneration(self, task_id: str) -> dict[str, Any]:
|
||
"""恢复旧版在真正开始生成前误删的上一轮结果。
|
||
|
||
旧实现会在 ``POST /regenerate`` 时立即把已发布任务置为 pending、
|
||
清空结果并解除输出指针。三个已发布数据集仍是独立完整产物，因此只在
|
||
这个特征完全匹配时，使用其记录恢复结果和任务状态。该操作幂等，不会
|
||
触碰正常的新建待生成任务或已经开始的新一轮生成。
|
||
"""
|
||
|
||
with self.connect() as conn:
|
||
task = self._task_in_connection(conn, task_id, for_update=True)
|
||
if (
|
||
task.get("status") != "pending"
|
||
or task.get("generation_run_id")
|
||
or task.get("output_dataset_id")
|
||
or int(task.get("output_count") or 0) != 0
|
||
):
|
||
return {"recovered": False, "result_count": 0}
|
||
|
||
result_count = int(
|
||
(
|
||
conn.execute(
|
||
"SELECT COUNT(*) AS count FROM data_process_results WHERE task_id=%s",
|
||
(task_id,),
|
||
).fetchone()
|
||
or {}
|
||
).get("count")
|
||
or 0
|
||
)
|
||
if result_count:
|
||
return {"recovered": False, "result_count": result_count}
|
||
|
||
datasets = conn.execute(
|
||
"""
|
||
SELECT id, type, count, created_at
|
||
FROM datasets
|
||
WHERE source='task' AND deleted_at IS NULL
|
||
AND (
|
||
source_task_id=%s
|
||
OR (source_task_id IS NULL AND task_id=%s)
|
||
)
|
||
ORDER BY CASE type
|
||
WHEN 'train' THEN 1 WHEN 'val' THEN 2 WHEN 'test' THEN 3 ELSE 4
|
||
END, created_at, id
|
||
""",
|
||
(task_id, task_id),
|
||
).fetchall()
|
||
train_dataset = next(
|
||
(dataset for dataset in datasets if dataset.get("type") == "train"),
|
||
None,
|
||
)
|
||
if not train_dataset:
|
||
return {"recovered": False, "result_count": 0}
|
||
|
||
dataset_ids = [str(dataset["id"]) for dataset in datasets]
|
||
records = conn.execute(
|
||
"""
|
||
SELECT id, dataset_id, line_no, split, instruction, input, output,
|
||
raw, status, source_result_id, preview_item_id, created_at
|
||
FROM dataset_records
|
||
WHERE dataset_id = ANY(%s)
|
||
ORDER BY created_at, dataset_id, line_no NULLS LAST, id
|
||
""",
|
||
(dataset_ids,),
|
||
).fetchall()
|
||
if not records:
|
||
return {"recovered": False, "result_count": 0}
|
||
|
||
preview_rows = conn.execute(
|
||
"SELECT id FROM data_process_preview_items WHERE task_id=%s",
|
||
(task_id,),
|
||
).fetchall()
|
||
preview_ids = {str(row["id"]) for row in preview_rows}
|
||
used_result_ids: set[str] = set()
|
||
recovered_count = 0
|
||
for record in records:
|
||
raw = _json_value(record.get("raw"), {})
|
||
raw = raw if isinstance(raw, dict) else {}
|
||
candidate_id = str(
|
||
record.get("source_result_id")
|
||
or raw.get("source_result_id")
|
||
or ""
|
||
)
|
||
result_id = (
|
||
candidate_id
|
||
if candidate_id and candidate_id not in used_result_ids
|
||
else new_id("dpr")
|
||
)
|
||
used_result_ids.add(result_id)
|
||
candidate_preview_id = str(
|
||
record.get("preview_item_id")
|
||
or raw.get("preview_item_id")
|
||
or ""
|
||
)
|
||
preview_item_id = (
|
||
candidate_preview_id if candidate_preview_id in preview_ids else None
|
||
)
|
||
instruction = str(record.get("instruction") or raw.get("instruction") or "")
|
||
input_text = str(record.get("input") or raw.get("input") or "")
|
||
output = str(record.get("output") or raw.get("output") or "")
|
||
split = str(record.get("split") or raw.get("split") or "") or None
|
||
status = str(record.get("status") or "valid")
|
||
if status not in {"valid", "modified", "invalid"}:
|
||
status = "valid"
|
||
created_at = record.get("created_at") or utcnow()
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO data_process_results
|
||
(id, task_id, preview_item_id, instruction, input, output,
|
||
original_instruction, original_input, original_output, status,
|
||
error, split, quality_score, created_at, updated_at)
|
||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s,
|
||
NULL, %s, '{}', %s, %s)
|
||
""",
|
||
(
|
||
result_id,
|
||
task_id,
|
||
preview_item_id,
|
||
instruction,
|
||
input_text,
|
||
output,
|
||
instruction,
|
||
input_text,
|
||
output,
|
||
status,
|
||
split,
|
||
created_at,
|
||
created_at,
|
||
),
|
||
)
|
||
conn.execute(
|
||
"""
|
||
UPDATE dataset_records
|
||
SET source_result_id=%s, preview_item_id=%s
|
||
WHERE id=%s
|
||
""",
|
||
(result_id, preview_item_id, record["id"]),
|
||
)
|
||
recovered_count += 1
|
||
|
||
now = utcnow()
|
||
conn.execute(
|
||
"""
|
||
UPDATE data_process_tasks
|
||
SET status='completed', progress=100, output_dataset_id=%s,
|
||
output_count=%s, filtered_count=0, duplicate_count=0,
|
||
error_count=(SELECT COUNT(*) FROM data_process_results
|
||
WHERE task_id=%s AND status='invalid'),
|
||
failure_reason=NULL, updated_at=%s
|
||
WHERE id=%s
|
||
""",
|
||
(
|
||
train_dataset["id"],
|
||
recovered_count,
|
||
task_id,
|
||
now,
|
||
task_id,
|
||
),
|
||
)
|
||
return {"recovered": True, "result_count": recovered_count}
|
||
|
||
def update_task(self, task_id: str, payload: dict[str, Any]) -> dict[str, Any]:
|
||
allowed = {
|
||
"name",
|
||
"description",
|
||
"process_type",
|
||
"source_dataset_id",
|
||
}
|
||
values: dict[str, Any] = {key: value for key, value in payload.items() if key in allowed}
|
||
if not values and payload.get("config") is None:
|
||
return self.get_task(task_id)
|
||
try:
|
||
with self.connect() as conn:
|
||
task = self._task_in_connection(conn, task_id, for_update=True)
|
||
self._ensure_editable(task)
|
||
regeneration_prepared = _is_regeneration_prepared(task)
|
||
if regeneration_prepared and any(
|
||
key in payload and payload.get(key) != task.get(key)
|
||
for key in ("process_type", "source_dataset_id")
|
||
):
|
||
raise InvalidStateError(
|
||
"process type and source dataset cannot change during regeneration"
|
||
)
|
||
if payload.get("config") is not None:
|
||
next_config = {
|
||
key: value
|
||
for key, value in payload["config"].items()
|
||
if key != _REGENERATION_MARKER_KEY
|
||
}
|
||
current_marker = _regeneration_marker(task)
|
||
if current_marker:
|
||
next_config[_REGENERATION_MARKER_KEY] = current_marker
|
||
values["config"] = json_dumps(next_config)
|
||
invalidates_results = (
|
||
("config" in payload and payload.get("config") != task.get("config"))
|
||
or (
|
||
"process_type" in payload
|
||
and payload.get("process_type") != task.get("process_type")
|
||
)
|
||
or (
|
||
"source_dataset_id" in payload
|
||
and payload.get("source_dataset_id") != task.get("source_dataset_id")
|
||
)
|
||
)
|
||
if invalidates_results and not regeneration_prepared:
|
||
values.update(
|
||
{
|
||
"status": "pending",
|
||
"progress": 0,
|
||
"output_count": 0,
|
||
"filtered_count": 0,
|
||
"duplicate_count": 0,
|
||
"error_count": 0,
|
||
"failure_reason": None,
|
||
"generation_run_id": None,
|
||
"started_at": None,
|
||
"completed_at": None,
|
||
}
|
||
)
|
||
conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
|
||
conn.execute(
|
||
"DELETE FROM data_process_preview_items WHERE task_id=%s", (task_id,)
|
||
)
|
||
if (
|
||
"process_type" in payload
|
||
and payload.get("process_type") != task.get("process_type")
|
||
):
|
||
conn.execute(
|
||
"""
|
||
UPDATE data_process_source_files
|
||
SET deleted_at=%s, updated_at=%s
|
||
WHERE task_id=%s AND deleted_at IS NULL
|
||
""",
|
||
(utcnow(), utcnow(), task_id),
|
||
)
|
||
values["input_count"] = 0
|
||
values["updated_at"] = utcnow()
|
||
assignments = ", ".join(f"{key}=%s" for key in values)
|
||
row = conn.execute(
|
||
f"UPDATE data_process_tasks SET {assignments} WHERE id=%s RETURNING *",
|
||
[*values.values(), task_id],
|
||
).fetchone()
|
||
except psycopg.errors.UniqueViolation as exc:
|
||
raise ConflictError("data process task name already exists") from exc
|
||
return _public_task(_decode_row(row)) or {}
|
||
|
||
def prepare_regeneration(
|
||
self,
|
||
task_id: str,
|
||
payload: dict[str, Any],
|
||
) -> dict[str, Any]:
|
||
"""非破坏性地保存重新生成配置。
|
||
|
||
准备阶段保留任务当前状态、结果、切片及已发布数据集。真正开始
|
||
生成时，才在同一事务内切换运行状态并清理上一轮结果。
|
||
"""
|
||
|
||
# 先修复曾被旧版 prepare 提前清空的任务，再建立新的非破坏性草稿标记。
|
||
self.recover_legacy_aborted_regeneration(task_id)
|
||
try:
|
||
with self.connect() as conn:
|
||
task = self._task_in_connection(conn, task_id, for_update=True)
|
||
if task["status"] == "running":
|
||
raise ConflictError("running task cannot be prepared for regeneration")
|
||
|
||
current_updated_at = _serialize_value(task.get("updated_at"))
|
||
if payload["expected_updated_at"] != current_updated_at:
|
||
raise ConflictError("data process task was modified by another request")
|
||
|
||
process_type = str(payload["process_type"])
|
||
if process_type != str(task["process_type"]):
|
||
raise InvalidStateError("process_type cannot be changed during regeneration")
|
||
|
||
current_config = dict(task.get("config") or {})
|
||
next_config = dict(payload.get("config") or {})
|
||
next_config.pop(_REGENERATION_MARKER_KEY, None)
|
||
preview_invalidated = _preview_config_changed(
|
||
process_type,
|
||
current_config,
|
||
next_config,
|
||
)
|
||
now = utcnow()
|
||
next_config[_REGENERATION_MARKER_KEY] = {
|
||
"prepared": True,
|
||
"preview_invalidated": preview_invalidated,
|
||
"prepared_at": now,
|
||
}
|
||
# 002 迁移前发布的数据集只有 task_id。先补齐新关联字段，保证
|
||
# 解除任务输出指针后，详情和后续重新发布仍能定位原来的三份数据集。
|
||
conn.execute(
|
||
"""
|
||
UPDATE datasets
|
||
SET source_task_id=%s, updated_at=%s
|
||
WHERE source='task' AND source_task_id IS NULL AND task_id=%s
|
||
AND deleted_at IS NULL
|
||
""",
|
||
(task_id, now, task_id),
|
||
)
|
||
published_row = conn.execute(
|
||
"""
|
||
SELECT EXISTS(
|
||
SELECT 1 FROM datasets
|
||
WHERE source='task' AND deleted_at IS NULL
|
||
AND (
|
||
source_task_id=%s
|
||
OR (source_task_id IS NULL AND task_id=%s)
|
||
)
|
||
) AS exists
|
||
""",
|
||
(task_id, task_id),
|
||
).fetchone()
|
||
published_outputs_preserved = bool(task.get("output_dataset_id")) or bool(
|
||
published_row and published_row.get("exists")
|
||
)
|
||
|
||
row = conn.execute(
|
||
"""
|
||
UPDATE data_process_tasks
|
||
SET name=%s, description=%s, config=%s, updated_at=%s
|
||
WHERE id=%s
|
||
RETURNING *
|
||
""",
|
||
(
|
||
payload["name"],
|
||
payload.get("description") or "",
|
||
json_dumps(next_config),
|
||
now,
|
||
task_id,
|
||
),
|
||
).fetchone()
|
||
except psycopg.errors.UniqueViolation as exc:
|
||
raise ConflictError("data process task name already exists") from exc
|
||
return {
|
||
"task": _public_task(_decode_row(row)) or {},
|
||
"preview_invalidated": preview_invalidated,
|
||
"published_outputs_preserved": published_outputs_preserved,
|
||
}
|
||
|
||
def delete_task(self, task_id: str, *, deleted_by: str | None = None) -> None:
|
||
with self.connect() as conn:
|
||
task = self._task_in_connection(conn, task_id, for_update=True)
|
||
if task["status"] == "running":
|
||
raise InvalidStateError("running task must be stopped before deletion")
|
||
now = utcnow()
|
||
conn.execute(
|
||
"""
|
||
UPDATE data_process_tasks
|
||
SET deleted_at=%s, deleted_by=%s, updated_at=%s
|
||
WHERE id=%s
|
||
""",
|
||
(now, deleted_by, now, task_id),
|
||
)
|
||
|
||
def list_source_files(self, task_id: str) -> list[dict[str, Any]]:
|
||
self.get_task(task_id)
|
||
with self.connect() as conn:
|
||
rows = conn.execute(
|
||
"""
|
||
SELECT id, task_id, storage_object_id, name, size_bytes, record_count,
|
||
file_format, checksum_sha256, version_no, content_preview, metadata,
|
||
tenant_id, project_id, created_by, created_at, updated_at
|
||
FROM data_process_source_files
|
||
WHERE task_id=%s AND deleted_at IS NULL
|
||
ORDER BY created_at, id
|
||
""",
|
||
(task_id,),
|
||
).fetchall()
|
||
return [_decode_row(row) or {} for row in rows]
|
||
|
||
def add_source_file(
|
||
self,
|
||
task_id: str,
|
||
*,
|
||
name: str,
|
||
content: str,
|
||
raw_size: int,
|
||
checksum_sha256: str,
|
||
file_format: str,
|
||
record_count: int,
|
||
metadata: dict[str, Any] | None = None,
|
||
created_by: str | None = None,
|
||
source_file_id: str | None = None,
|
||
storage_object_id: str | None = None,
|
||
) -> dict[str, Any]:
|
||
return self.add_source_files(
|
||
task_id,
|
||
[
|
||
{
|
||
"name": name,
|
||
"content": content,
|
||
"raw_size": raw_size,
|
||
"checksum_sha256": checksum_sha256,
|
||
"file_format": file_format,
|
||
"record_count": record_count,
|
||
"metadata": metadata or {},
|
||
"created_by": created_by,
|
||
"id": source_file_id,
|
||
"storage_object_id": storage_object_id,
|
||
}
|
||
],
|
||
)[0]
|
||
|
||
def add_source_files(
|
||
self,
|
||
task_id: str,
|
||
files: Sequence[dict[str, Any]],
|
||
) -> list[dict[str, Any]]:
|
||
"""在同一事务中登记一个上传批次，任一文件失败则全部回滚。"""
|
||
|
||
if not files:
|
||
raise DataProcessStoreError("at least one source file is required")
|
||
now = utcnow()
|
||
created: list[dict[str, Any]] = []
|
||
try:
|
||
with self.connect() as conn:
|
||
task = self._task_in_connection(conn, task_id, for_update=True)
|
||
self._ensure_editable(task)
|
||
for payload in files:
|
||
file_id = str(payload.get("id") or new_id("dpsf"))
|
||
storage_object_id, metadata_payload = _source_storage_descriptor(
|
||
payload,
|
||
task_id,
|
||
file_id,
|
||
)
|
||
row = conn.execute(
|
||
"""
|
||
INSERT INTO data_process_source_files
|
||
(id, task_id, storage_object_id, name, size_bytes, record_count,
|
||
file_format, checksum_sha256, version_no, content, content_preview,
|
||
metadata, tenant_id, project_id,
|
||
created_by, created_at, updated_at)
|
||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, 1, %s, %s, %s, %s,
|
||
%s, %s, %s, %s)
|
||
RETURNING id, task_id, storage_object_id, name, size_bytes, record_count,
|
||
file_format, checksum_sha256, version_no, content_preview, metadata,
|
||
tenant_id, project_id, created_by, created_at, updated_at
|
||
""",
|
||
(
|
||
file_id,
|
||
task_id,
|
||
storage_object_id,
|
||
payload["name"],
|
||
payload["raw_size"],
|
||
payload["record_count"],
|
||
payload["file_format"],
|
||
payload["checksum_sha256"],
|
||
payload["content"],
|
||
str(payload["content"])[:2000],
|
||
json_dumps(metadata_payload),
|
||
task.get("tenant_id"),
|
||
task.get("project_id"),
|
||
payload.get("created_by") or task.get("created_by"),
|
||
now,
|
||
now,
|
||
),
|
||
).fetchone()
|
||
created.append(_decode_row(row) or {})
|
||
preview_row = conn.execute(
|
||
"""
|
||
SELECT COUNT(*) AS count FROM data_process_preview_items
|
||
WHERE task_id=%s
|
||
""",
|
||
(task_id,),
|
||
).fetchone()
|
||
preview_count = int((preview_row or {}).get("count") or 0)
|
||
if _is_regeneration_prepared(task):
|
||
conn.execute(
|
||
"""
|
||
UPDATE data_process_tasks
|
||
SET input_count=(
|
||
SELECT COALESCE(SUM(record_count), 0)
|
||
FROM data_process_source_files
|
||
WHERE task_id=%s AND deleted_at IS NULL
|
||
), updated_at=%s
|
||
WHERE id=%s
|
||
""",
|
||
(task_id, now, task_id),
|
||
)
|
||
else:
|
||
# 前端会只对本次新增的源文件构建预览，因此保留旧文件切片，
|
||
# 但普通未发布任务的旧生成结果已经不再有效。
|
||
conn.execute(
|
||
"DELETE FROM data_process_results WHERE task_id=%s", (task_id,)
|
||
)
|
||
conn.execute(
|
||
"""
|
||
UPDATE data_process_tasks
|
||
SET status='pending', progress=%s, output_count=0, filtered_count=0,
|
||
duplicate_count=0, error_count=0, failure_reason=NULL,
|
||
generation_run_id=NULL, started_at=NULL, completed_at=NULL,
|
||
input_count=(
|
||
SELECT COALESCE(SUM(record_count), 0)
|
||
FROM data_process_source_files
|
||
WHERE task_id=%s AND deleted_at IS NULL
|
||
), updated_at=%s
|
||
WHERE id=%s
|
||
""",
|
||
(20 if preview_count else 0, task_id, now, task_id),
|
||
)
|
||
except psycopg.errors.UniqueViolation as exc:
|
||
raise ConflictError(
|
||
"the same source file content is already attached to this task"
|
||
) from exc
|
||
return created
|
||
|
||
def get_source_file(
|
||
self, task_id: str, file_id: str, *, include_content: bool = True
|
||
) -> dict[str, Any]:
|
||
# 先验证父任务仍然可见，避免软删除任务后通过已知文件 ID 读取正文。
|
||
self.get_task(task_id)
|
||
content_column = ", content" if include_content else ""
|
||
with self.connect() as conn:
|
||
row = conn.execute(
|
||
f"""
|
||
SELECT id, task_id, storage_object_id, name, size_bytes, record_count,
|
||
file_format, checksum_sha256, version_no, content_preview, metadata,
|
||
tenant_id, project_id, created_by, created_at, updated_at{content_column}
|
||
FROM data_process_source_files
|
||
WHERE id=%s AND task_id=%s AND deleted_at IS NULL
|
||
""",
|
||
(file_id, task_id),
|
||
).fetchone()
|
||
if not row:
|
||
raise NotFoundError("source file not found")
|
||
return _decode_row(row) or {}
|
||
|
||
def source_content_window(
|
||
self, task_id: str, file_id: str, offset: int, limit: int
|
||
) -> dict[str, Any]:
|
||
source_file = self.get_source_file(task_id, file_id, include_content=True)
|
||
content = str(source_file.pop("content", ""))
|
||
window = content[offset : offset + limit]
|
||
return {
|
||
"file": source_file,
|
||
"content": window,
|
||
"offset": offset,
|
||
"limit": limit,
|
||
"total_chars": len(content),
|
||
"has_more": offset + len(window) < len(content),
|
||
}
|
||
|
||
def source_content_lines(
|
||
self,
|
||
task_id: str,
|
||
file_id: str,
|
||
start_line: int,
|
||
line_count: int,
|
||
) -> dict[str, Any]:
|
||
source_file = self.get_source_file(task_id, file_id, include_content=True)
|
||
content = str(source_file.pop("content", ""))
|
||
lines = content.splitlines(keepends=True)
|
||
start_index = min(len(lines), start_line - 1)
|
||
selected = lines[start_index : start_index + line_count]
|
||
end_line = start_index + len(selected)
|
||
return {
|
||
"file": source_file,
|
||
"content": "".join(selected),
|
||
"start_line": start_line,
|
||
"end_line": end_line,
|
||
"line_count": len(selected),
|
||
"total_lines": len(lines),
|
||
"has_more": end_line < len(lines),
|
||
}
|
||
|
||
def delete_source_file(self, task_id: str, file_id: str) -> None:
|
||
with self.connect() as conn:
|
||
task = self._task_in_connection(conn, task_id, for_update=True)
|
||
self._ensure_editable(task)
|
||
row = conn.execute(
|
||
"""
|
||
UPDATE data_process_source_files
|
||
SET deleted_at=%s, updated_at=%s
|
||
WHERE id=%s AND task_id=%s AND deleted_at IS NULL
|
||
RETURNING id
|
||
""",
|
||
(utcnow(), utcnow(), file_id, task_id),
|
||
).fetchone()
|
||
if not row:
|
||
raise NotFoundError("source file not found")
|
||
conn.execute(
|
||
"DELETE FROM data_process_preview_items WHERE source_file_id=%s", (file_id,)
|
||
)
|
||
now = utcnow()
|
||
if _is_regeneration_prepared(task):
|
||
conn.execute(
|
||
"""
|
||
UPDATE data_process_tasks
|
||
SET input_count=(SELECT COALESCE(SUM(record_count), 0)
|
||
FROM data_process_source_files
|
||
WHERE task_id=%s AND deleted_at IS NULL),
|
||
updated_at=%s
|
||
WHERE id=%s
|
||
""",
|
||
(task_id, now, task_id),
|
||
)
|
||
else:
|
||
conn.execute(
|
||
"DELETE FROM data_process_results WHERE task_id=%s", (task_id,)
|
||
)
|
||
conn.execute(
|
||
"""
|
||
UPDATE data_process_tasks
|
||
SET status='pending', progress=0, output_count=0, filtered_count=0,
|
||
duplicate_count=0, error_count=0, failure_reason=NULL,
|
||
input_count=(SELECT COALESCE(SUM(record_count), 0)
|
||
FROM data_process_source_files
|
||
WHERE task_id=%s AND deleted_at IS NULL),
|
||
updated_at=%s
|
||
WHERE id=%s
|
||
""",
|
||
(task_id, now, task_id),
|
||
)
|
||
|
||
def replace_preview_items(
|
||
self,
|
||
task_id: str,
|
||
items: Sequence[dict[str, Any]],
|
||
*,
|
||
source_file_ids: Sequence[str] | None = None,
|
||
) -> list[dict[str, Any]]:
|
||
selected_ids = (
|
||
list(dict.fromkeys(str(file_id) for file_id in source_file_ids))
|
||
if source_file_ids is not None
|
||
else None
|
||
)
|
||
if selected_ids is not None:
|
||
if not selected_ids or any(not file_id for file_id in selected_ids):
|
||
raise ValueError("source_file_ids must contain non-empty ids")
|
||
selected_set = set(selected_ids)
|
||
unexpected = {
|
||
str(item.get("source_file_id") or "")
|
||
for item in items
|
||
if str(item.get("source_file_id") or "") not in selected_set
|
||
}
|
||
if unexpected:
|
||
raise ValueError("preview items contain an unselected source file")
|
||
|
||
now = utcnow()
|
||
with self.connect() as conn:
|
||
task = self._task_in_connection(conn, task_id, for_update=True)
|
||
self._ensure_editable(task)
|
||
regeneration_prepared = _is_regeneration_prepared(task)
|
||
if not regeneration_prepared:
|
||
conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
|
||
if selected_ids is None:
|
||
conn.execute(
|
||
"DELETE FROM data_process_preview_items WHERE task_id=%s", (task_id,)
|
||
)
|
||
else:
|
||
rows = conn.execute(
|
||
"""
|
||
SELECT id FROM data_process_source_files
|
||
WHERE task_id=%s AND deleted_at IS NULL AND id=ANY(%s)
|
||
""",
|
||
(task_id, selected_ids),
|
||
).fetchall()
|
||
found = {str(row["id"]) for row in rows}
|
||
missing = set(selected_ids) - found
|
||
if missing:
|
||
raise NotFoundError(
|
||
f"source files not found: {', '.join(sorted(missing))}"
|
||
)
|
||
conn.execute(
|
||
"""
|
||
DELETE FROM data_process_preview_items
|
||
WHERE task_id=%s AND source_file_id=ANY(%s)
|
||
""",
|
||
(task_id, selected_ids),
|
||
)
|
||
created: list[dict[str, Any]] = []
|
||
for item in items:
|
||
row = conn.execute(
|
||
"""
|
||
INSERT INTO data_process_preview_items
|
||
(id, task_id, source_file_id, original_content, edited_content,
|
||
source_start, source_end, source_start_line, source_end_line,
|
||
token_count, status, quality_score, created_at, updated_at)
|
||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
|
||
RETURNING *
|
||
""",
|
||
(
|
||
item.get("id") or new_id("dpp"),
|
||
task_id,
|
||
item.get("source_file_id"),
|
||
item.get("original_content") or "",
|
||
item.get("edited_content", item.get("original_content") or ""),
|
||
item.get("source_start"),
|
||
item.get("source_end"),
|
||
item.get("source_start_line"),
|
||
item.get("source_end_line"),
|
||
max(0, int(item.get("token_count") or 0)),
|
||
item.get("status") or "original",
|
||
json_dumps(item.get("quality_score") or {}),
|
||
now,
|
||
now,
|
||
),
|
||
).fetchone()
|
||
created.append(_decode_row(row) or {})
|
||
if regeneration_prepared:
|
||
conn.execute(
|
||
"UPDATE data_process_tasks SET updated_at=%s WHERE id=%s",
|
||
(now, task_id),
|
||
)
|
||
else:
|
||
conn.execute(
|
||
"""
|
||
UPDATE data_process_tasks
|
||
SET status='pending', progress=20, output_count=0, filtered_count=0,
|
||
duplicate_count=0, error_count=0, failure_reason=NULL, updated_at=%s
|
||
WHERE id=%s
|
||
""",
|
||
(now, task_id),
|
||
)
|
||
return created
|
||
|
||
def list_preview_items(
|
||
self,
|
||
task_id: str,
|
||
*,
|
||
source_file_id: str | None = None,
|
||
page: int = 1,
|
||
page_size: int = 200,
|
||
keyword: str | None = None,
|
||
) -> dict[str, Any]:
|
||
self.get_task(task_id)
|
||
clauses = ["task_id=%s"]
|
||
params: list[Any] = [task_id]
|
||
if source_file_id:
|
||
clauses.append("source_file_id=%s")
|
||
params.append(source_file_id)
|
||
if keyword:
|
||
clauses.append("(original_content ILIKE %s OR edited_content ILIKE %s)")
|
||
pattern = f"%{keyword.strip()}%"
|
||
params.extend([pattern, pattern])
|
||
where = " AND ".join(clauses)
|
||
with self.connect() as conn:
|
||
total = conn.execute(
|
||
f"SELECT COUNT(*) AS count FROM data_process_preview_items WHERE {where}", params
|
||
).fetchone()["count"]
|
||
rows = conn.execute(
|
||
f"""
|
||
SELECT * FROM data_process_preview_items
|
||
WHERE {where}
|
||
ORDER BY source_file_id NULLS LAST, source_start NULLS LAST, created_at, id
|
||
LIMIT %s OFFSET %s
|
||
""",
|
||
[*params, page_size, (page - 1) * page_size],
|
||
).fetchall()
|
||
return {
|
||
"items": [_decode_row(row) for row in rows],
|
||
"total": int(total),
|
||
"page": page,
|
||
"page_size": page_size,
|
||
}
|
||
|
||
def get_preview_item(self, task_id: str, preview_id: str) -> dict[str, Any]:
|
||
with self.connect() as conn:
|
||
row = conn.execute(
|
||
"SELECT * FROM data_process_preview_items WHERE id=%s AND task_id=%s",
|
||
(preview_id, task_id),
|
||
).fetchone()
|
||
if not row:
|
||
raise NotFoundError("preview item not found")
|
||
return _decode_row(row) or {}
|
||
|
||
def create_preview_item(self, task_id: str, item: dict[str, Any]) -> dict[str, Any]:
|
||
now = utcnow()
|
||
with self.connect() as conn:
|
||
task = self._task_in_connection(conn, task_id, for_update=True)
|
||
self._ensure_editable(task)
|
||
if item.get("source_file_id"):
|
||
source = conn.execute(
|
||
"""
|
||
SELECT id FROM data_process_source_files
|
||
WHERE id=%s AND task_id=%s AND deleted_at IS NULL
|
||
""",
|
||
(item["source_file_id"], task_id),
|
||
).fetchone()
|
||
if not source:
|
||
raise NotFoundError("source file not found")
|
||
row = conn.execute(
|
||
"""
|
||
INSERT INTO data_process_preview_items
|
||
(id, task_id, source_file_id, original_content, edited_content,
|
||
source_start, source_end, source_start_line, source_end_line,
|
||
token_count, status, quality_score, created_at, updated_at)
|
||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
|
||
RETURNING *
|
||
""",
|
||
(
|
||
new_id("dpp"),
|
||
task_id,
|
||
item.get("source_file_id"),
|
||
item.get("original_content") or "",
|
||
item.get("edited_content") or "",
|
||
item.get("source_start"),
|
||
item.get("source_end"),
|
||
item.get("source_start_line"),
|
||
item.get("source_end_line"),
|
||
max(0, int(item.get("token_count") or 0)),
|
||
item.get("status") or "manual",
|
||
json_dumps(item.get("quality_score") or {}),
|
||
now,
|
||
now,
|
||
),
|
||
).fetchone()
|
||
self._invalidate_results(conn, task, task_id, now)
|
||
return _decode_row(row) or {}
|
||
|
||
def update_preview_item(
|
||
self, task_id: str, preview_id: str, payload: dict[str, Any]
|
||
) -> dict[str, Any]:
|
||
with self.connect() as conn:
|
||
task = self._task_in_connection(conn, task_id, for_update=True)
|
||
self._ensure_editable(task)
|
||
existing = conn.execute(
|
||
"SELECT * FROM data_process_preview_items WHERE id=%s AND task_id=%s",
|
||
(preview_id, task_id),
|
||
).fetchone()
|
||
if not existing:
|
||
raise NotFoundError("preview item not found")
|
||
expected_updated_at = payload.get("expected_updated_at")
|
||
current_updated_at = _serialize_value(existing.get("updated_at"))
|
||
if expected_updated_at and expected_updated_at != current_updated_at:
|
||
raise ConflictError("preview item was modified by another request")
|
||
edited = payload["edited_content"]
|
||
status = payload.get("status")
|
||
if not status:
|
||
if not edited.strip():
|
||
status = "invalid"
|
||
elif edited == existing["original_content"]:
|
||
status = "original"
|
||
else:
|
||
status = "modified"
|
||
now = utcnow()
|
||
row = conn.execute(
|
||
"""
|
||
UPDATE data_process_preview_items
|
||
SET edited_content=%s, token_count=%s, status=%s, quality_score=%s,
|
||
updated_at=%s
|
||
WHERE id=%s AND task_id=%s
|
||
RETURNING *
|
||
""",
|
||
(
|
||
edited,
|
||
estimate_token_count(edited),
|
||
status,
|
||
json_dumps(payload.get("quality_score") or {}),
|
||
now,
|
||
preview_id,
|
||
task_id,
|
||
),
|
||
).fetchone()
|
||
self._invalidate_results(conn, task, task_id, now)
|
||
return _decode_row(row) or {}
|
||
|
||
def delete_preview_item(self, task_id: str, preview_id: str) -> None:
|
||
with self.connect() as conn:
|
||
task = self._task_in_connection(conn, task_id, for_update=True)
|
||
self._ensure_editable(task)
|
||
row = conn.execute(
|
||
"DELETE FROM data_process_preview_items WHERE id=%s AND task_id=%s RETURNING id",
|
||
(preview_id, task_id),
|
||
).fetchone()
|
||
if not row:
|
||
raise NotFoundError("preview item not found")
|
||
self._invalidate_results(conn, task, task_id, utcnow())
|
||
|
||
def _invalidate_results(
|
||
self,
|
||
conn: psycopg.Connection[dict[str, Any]],
|
||
task: dict[str, Any],
|
||
task_id: str,
|
||
now: str,
|
||
) -> None:
|
||
if _is_regeneration_prepared(task):
|
||
conn.execute(
|
||
"UPDATE data_process_tasks SET updated_at=%s WHERE id=%s",
|
||
(now, task_id),
|
||
)
|
||
return
|
||
conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
|
||
conn.execute(
|
||
"""
|
||
UPDATE data_process_tasks
|
||
SET status='pending', progress=20, output_count=0, filtered_count=0,
|
||
duplicate_count=0, error_count=0, failure_reason=NULL,
|
||
generation_run_id=NULL, updated_at=%s
|
||
WHERE id=%s
|
||
""",
|
||
(now, task_id),
|
||
)
|
||
|
||
def start_generation(self, task_id: str, *, replace_existing: bool = True) -> dict[str, Any]:
|
||
if not replace_existing:
|
||
raise DataProcessStoreError("incremental generation is not supported")
|
||
with self.connect() as conn:
|
||
task = self._task_in_connection(conn, task_id, for_update=True)
|
||
regeneration_prepared = _is_regeneration_prepared(task)
|
||
if task.get("output_dataset_id") and not regeneration_prepared:
|
||
raise InvalidStateError("published task cannot be regenerated")
|
||
if task["status"] == "running":
|
||
raise ConflictError("data process task is already running")
|
||
preview_count = conn.execute(
|
||
"SELECT COUNT(*) AS count FROM data_process_preview_items WHERE task_id=%s",
|
||
(task_id,),
|
||
).fetchone()["count"]
|
||
if not preview_count:
|
||
raise InvalidStateError("preview must be built before generation")
|
||
conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
|
||
now = utcnow()
|
||
generation_run_id = new_id("dprun")
|
||
next_config = dict(task.get("config") or {})
|
||
next_config.pop(_REGENERATION_MARKER_KEY, None)
|
||
row = conn.execute(
|
||
"""
|
||
UPDATE data_process_tasks
|
||
SET config=%s, status='running', progress=30, failure_reason=NULL,
|
||
started_at=%s, completed_at=NULL, output_dataset_id=NULL,
|
||
output_count=0, filtered_count=0, duplicate_count=0, error_count=0,
|
||
generation_run_id=%s, updated_at=%s
|
||
WHERE id=%s
|
||
RETURNING *
|
||
""",
|
||
(json_dumps(next_config), now, generation_run_id, now, task_id),
|
||
).fetchone()
|
||
return _decode_row(row) or {}
|
||
|
||
def stop_task(self, task_id: str) -> dict[str, Any]:
|
||
with self.connect() as conn:
|
||
task = self._task_in_connection(conn, task_id, for_update=True)
|
||
if task["status"] != "running":
|
||
raise InvalidStateError("only a running task can be stopped")
|
||
now = utcnow()
|
||
row = conn.execute(
|
||
"""
|
||
UPDATE data_process_tasks
|
||
SET status='stopped', failure_reason=NULL, generation_run_id=NULL,
|
||
updated_at=%s
|
||
WHERE id=%s RETURNING *
|
||
""",
|
||
(now, task_id),
|
||
).fetchone()
|
||
return _decode_row(row) or {}
|
||
|
||
def generation_is_running(self, task_id: str, generation_run_id: str) -> bool:
|
||
task = self.get_task(task_id)
|
||
return (
|
||
task["status"] == "running"
|
||
and task.get("generation_run_id") == generation_run_id
|
||
)
|
||
|
||
def update_generation_progress(
|
||
self,
|
||
task_id: str,
|
||
generation_run_id: str,
|
||
processed_count: int,
|
||
total_count: int,
|
||
) -> bool:
|
||
ratio = processed_count / max(1, total_count)
|
||
progress = min(95.0, 30.0 + ratio * 65.0)
|
||
with self.connect() as conn:
|
||
row = conn.execute(
|
||
"""
|
||
UPDATE data_process_tasks
|
||
SET progress=%s, updated_at=%s
|
||
WHERE id=%s AND status='running' AND generation_run_id=%s
|
||
RETURNING id
|
||
""",
|
||
(progress, utcnow(), task_id, generation_run_id),
|
||
).fetchone()
|
||
return row is not None
|
||
|
||
def complete_generation(
|
||
self,
|
||
task_id: str,
|
||
results: Sequence[dict[str, Any]],
|
||
*,
|
||
generation_run_id: str,
|
||
filtered_count: int = 0,
|
||
duplicate_count: int = 0,
|
||
error_count: int = 0,
|
||
) -> dict[str, Any]:
|
||
now = utcnow()
|
||
with self.connect() as conn:
|
||
task = self._task_in_connection(conn, task_id, for_update=True)
|
||
if (
|
||
task["status"] != "running"
|
||
or task.get("generation_run_id") != generation_run_id
|
||
):
|
||
return task
|
||
conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
|
||
for result in results:
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO data_process_results
|
||
(id, task_id, preview_item_id, instruction, input, output,
|
||
original_instruction, original_input, original_output, status, error,
|
||
split, quality_score, created_at, updated_at)
|
||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
|
||
""",
|
||
(
|
||
result.get("id") or new_id("dpr"),
|
||
task_id,
|
||
result.get("preview_item_id"),
|
||
result.get("instruction") or "",
|
||
result.get("input") or "",
|
||
result.get("output") or "",
|
||
result.get("original_instruction", result.get("instruction") or ""),
|
||
result.get("original_input", result.get("input") or ""),
|
||
result.get("original_output", result.get("output") or ""),
|
||
result.get("status") or "valid",
|
||
result.get("error"),
|
||
result.get("split"),
|
||
json_dumps(result.get("quality_score") or {}),
|
||
now,
|
||
now,
|
||
),
|
||
)
|
||
row = conn.execute(
|
||
"""
|
||
UPDATE data_process_tasks
|
||
SET status='completed', progress=100, output_count=%s, filtered_count=%s,
|
||
duplicate_count=%s, error_count=%s, failure_reason=NULL,
|
||
completed_at=%s, generation_run_id=NULL, updated_at=%s
|
||
WHERE id=%s AND generation_run_id=%s RETURNING *
|
||
""",
|
||
(
|
||
len(results),
|
||
filtered_count,
|
||
duplicate_count,
|
||
error_count,
|
||
now,
|
||
now,
|
||
task_id,
|
||
generation_run_id,
|
||
),
|
||
).fetchone()
|
||
return _decode_row(row) or {}
|
||
|
||
def mark_failed(
|
||
self, task_id: str, reason: str, *, generation_run_id: str
|
||
) -> dict[str, Any]:
|
||
with self.connect() as conn:
|
||
task = self._task_in_connection(conn, task_id, for_update=True)
|
||
if (
|
||
task["status"] != "running"
|
||
or task.get("generation_run_id") != generation_run_id
|
||
):
|
||
return task
|
||
now = utcnow()
|
||
row = conn.execute(
|
||
"""
|
||
UPDATE data_process_tasks
|
||
SET status='failed', failure_reason=%s, completed_at=%s,
|
||
generation_run_id=NULL, updated_at=%s
|
||
WHERE id=%s AND generation_run_id=%s RETURNING *
|
||
""",
|
||
(reason[:4000], now, now, task_id, generation_run_id),
|
||
).fetchone()
|
||
return _decode_row(row) or {}
|
||
|
||
def progress(self, task_id: str) -> dict[str, Any]:
|
||
task = self.get_task(task_id)
|
||
return {
|
||
"task_id": task["id"],
|
||
"status": task["status"],
|
||
"progress": float(task.get("progress") or 0),
|
||
"input_count": int(task.get("input_count") or 0),
|
||
"output_count": int(task.get("output_count") or 0),
|
||
"filtered_count": int(task.get("filtered_count") or 0),
|
||
"duplicate_count": int(task.get("duplicate_count") or 0),
|
||
"error_count": int(task.get("error_count") or 0),
|
||
"failure_reason": task.get("failure_reason"),
|
||
"started_at": task.get("started_at"),
|
||
"completed_at": task.get("completed_at"),
|
||
}
|
||
|
||
def list_results(
|
||
self,
|
||
task_id: str,
|
||
*,
|
||
page: int = 1,
|
||
page_size: int = 100,
|
||
status: str | None = None,
|
||
split: str | None = None,
|
||
keyword: str | None = None,
|
||
) -> dict[str, Any]:
|
||
self.get_task(task_id)
|
||
clauses = ["task_id=%s"]
|
||
params: list[Any] = [task_id]
|
||
if status:
|
||
clauses.append("status=%s")
|
||
params.append(status)
|
||
if split:
|
||
clauses.append("split=%s")
|
||
params.append(split)
|
||
if keyword:
|
||
clauses.append("(instruction ILIKE %s OR input ILIKE %s OR output ILIKE %s)")
|
||
pattern = f"%{keyword.strip()}%"
|
||
params.extend([pattern, pattern, pattern])
|
||
where = " AND ".join(clauses)
|
||
with self.connect() as conn:
|
||
total = conn.execute(
|
||
f"SELECT COUNT(*) AS count FROM data_process_results WHERE {where}", params
|
||
).fetchone()["count"]
|
||
rows = conn.execute(
|
||
f"""
|
||
SELECT * FROM data_process_results WHERE {where}
|
||
ORDER BY created_at, id LIMIT %s OFFSET %s
|
||
""",
|
||
[*params, page_size, (page - 1) * page_size],
|
||
).fetchall()
|
||
return {
|
||
"items": [_decode_row(row) for row in rows],
|
||
"total": int(total),
|
||
"page": page,
|
||
"page_size": page_size,
|
||
}
|
||
|
||
def get_result(self, task_id: str, result_id: str) -> dict[str, Any]:
|
||
with self.connect() as conn:
|
||
row = conn.execute(
|
||
"SELECT * FROM data_process_results WHERE id=%s AND task_id=%s",
|
||
(result_id, task_id),
|
||
).fetchone()
|
||
if not row:
|
||
raise NotFoundError("data process result not found")
|
||
return _decode_row(row) or {}
|
||
|
||
def update_result(
|
||
self, task_id: str, result_id: str, payload: dict[str, Any]
|
||
) -> dict[str, Any]:
|
||
allowed = {"instruction", "input", "output", "quality_score"}
|
||
values = {key: value for key, value in payload.items() if key in allowed}
|
||
if "quality_score" in values:
|
||
values["quality_score"] = json_dumps(values["quality_score"])
|
||
if not values:
|
||
raise DataProcessStoreError("no result fields supplied")
|
||
with self.connect() as conn:
|
||
task = self._task_in_connection(conn, task_id, for_update=True)
|
||
if task["status"] == "running":
|
||
raise InvalidStateError("results cannot be edited while generation is running")
|
||
if task.get("output_dataset_id"):
|
||
raise InvalidStateError("published results cannot be edited")
|
||
current = conn.execute(
|
||
"SELECT * FROM data_process_results WHERE id=%s AND task_id=%s",
|
||
(result_id, task_id),
|
||
).fetchone()
|
||
if not current:
|
||
raise NotFoundError("data process result not found")
|
||
expected_updated_at = payload.get("expected_updated_at")
|
||
current_updated_at = _serialize_value(current.get("updated_at"))
|
||
if expected_updated_at and expected_updated_at != current_updated_at:
|
||
raise ConflictError("data process result was modified by another request")
|
||
merged = {**current, **values}
|
||
quality = payload.get("quality_score") or {}
|
||
instruction_valid = bool(str(merged.get("instruction") or "").strip())
|
||
output_valid = bool(str(merged.get("output") or "").strip())
|
||
reasoning_valid = (
|
||
_task_output_type(task) != "reasoning"
|
||
or _reasoning_output_is_valid(merged.get("output"))
|
||
)
|
||
hard_valid = instruction_valid and output_valid and reasoning_valid
|
||
quality_valid = bool(quality.get("is_valid", hard_valid))
|
||
changed = any(
|
||
str(merged.get(field) or "")
|
||
!= str(merged.get(f"original_{field}") or "")
|
||
for field in ("instruction", "input", "output")
|
||
)
|
||
status = "invalid" if not hard_valid or not quality_valid else (
|
||
"modified" if changed else "valid"
|
||
)
|
||
values["status"] = status
|
||
flags = quality.get("flags") if isinstance(quality, dict) else None
|
||
format_error = (
|
||
"思维链输出必须包含非空的 <think>...</think> 推理过程和最终答案"
|
||
if instruction_valid and output_valid and not reasoning_valid
|
||
else "Instruction 和 Output 不能为空"
|
||
if not instruction_valid or not output_valid
|
||
else None
|
||
)
|
||
values["error"] = ", ".join(str(flag) for flag in flags or []) or (
|
||
format_error
|
||
or ("quality validation failed" if status == "invalid" else None)
|
||
)
|
||
values["updated_at"] = utcnow()
|
||
assignments = ", ".join(f"{key}=%s" for key in values)
|
||
row = conn.execute(
|
||
f"""UPDATE data_process_results SET {assignments}
|
||
WHERE id=%s AND task_id=%s RETURNING *""",
|
||
[*values.values(), result_id, task_id],
|
||
).fetchone()
|
||
conn.execute(
|
||
"""
|
||
UPDATE data_process_tasks
|
||
SET error_count=(
|
||
SELECT COUNT(*) FROM data_process_results
|
||
WHERE task_id=%s AND status='invalid'
|
||
), updated_at=%s
|
||
WHERE id=%s
|
||
""",
|
||
(task_id, utcnow(), task_id),
|
||
)
|
||
return _decode_row(row) or {}
|
||
|
||
def get_generation_model(self, model_id: str) -> dict[str, Any]:
|
||
with self.connect() as conn:
|
||
row = conn.execute(
|
||
"""
|
||
SELECT id, name, type, purpose, model_source, description, path,
|
||
api_url, api_key, online_model_name, create_time
|
||
FROM models WHERE id=%s
|
||
""",
|
||
(model_id,),
|
||
).fetchone()
|
||
if not row:
|
||
raise NotFoundError("generation model not found")
|
||
return _decode_row(row) or {}
|
||
|
||
def save_generation_model_snapshot(
|
||
self,
|
||
task_id: str,
|
||
model_snapshot: dict[str, Any],
|
||
*,
|
||
generation_run_id: str,
|
||
) -> dict[str, Any]:
|
||
# API 密钥仅用于本次调用，绝不能进入任务配置、详情响应或审计快照。
|
||
safe_snapshot = {
|
||
key: value for key, value in model_snapshot.items() if key != "api_key"
|
||
}
|
||
with self.connect() as conn:
|
||
task = self._task_in_connection(conn, task_id, for_update=True)
|
||
if (
|
||
task["status"] != "running"
|
||
or task.get("generation_run_id") != generation_run_id
|
||
):
|
||
raise InvalidStateError("generation run is no longer active")
|
||
config = dict(task.get("config") or {})
|
||
config["generation_model_snapshot"] = safe_snapshot
|
||
row = conn.execute(
|
||
"""
|
||
UPDATE data_process_tasks SET config=%s, updated_at=%s
|
||
WHERE id=%s AND generation_run_id=%s RETURNING *
|
||
""",
|
||
(json_dumps(config), utcnow(), task_id, generation_run_id),
|
||
).fetchone()
|
||
return _decode_row(row) or {}
|
||
|
||
def publish(self, task_id: str, payload: dict[str, Any]) -> dict[str, Any]:
|
||
"""按精确配额发布训练、验证、测试三个独立数据集。"""
|
||
with self.connect() as conn:
|
||
task = self._task_in_connection(conn, task_id, for_update=True)
|
||
if _is_regeneration_prepared(task):
|
||
raise InvalidStateError(
|
||
"regeneration must start and complete before publishing"
|
||
)
|
||
if task["status"] != "completed":
|
||
raise InvalidStateError("only a completed task can be published")
|
||
rows = conn.execute(
|
||
"""
|
||
SELECT * FROM data_process_results
|
||
WHERE task_id=%s ORDER BY created_at, id
|
||
""",
|
||
(task_id,),
|
||
).fetchall()
|
||
if not rows:
|
||
raise InvalidStateError("task has no results to publish")
|
||
invalid_count = sum(
|
||
1
|
||
for row in rows
|
||
if row["status"] == "invalid"
|
||
or not str(row.get("instruction") or "").strip()
|
||
or not str(row.get("output") or "").strip()
|
||
or (
|
||
_task_output_type(task) == "reasoning"
|
||
and not _reasoning_output_is_valid(row.get("output"))
|
||
)
|
||
)
|
||
if invalid_count:
|
||
raise InvalidStateError(f"task contains {invalid_count} invalid results")
|
||
|
||
now = utcnow()
|
||
requested_split = payload.get("split") or {
|
||
"train": 80,
|
||
"validation": 10,
|
||
"test": 10,
|
||
}
|
||
assignments = stable_split_assignments(
|
||
[str(row["id"]) for row in rows],
|
||
requested_split,
|
||
seed=task_id,
|
||
)
|
||
records = [
|
||
{
|
||
"instruction": row["instruction"],
|
||
"input": row["input"],
|
||
"output": row["output"],
|
||
"split": assignment,
|
||
}
|
||
for row, assignment in zip(rows, assignments, strict=True)
|
||
]
|
||
split_order = ("train", "validation", "test")
|
||
split_counts = {
|
||
split_name: assignments.count(split_name) for split_name in split_order
|
||
}
|
||
split_specs: list[dict[str, Any]] = []
|
||
for split_name in split_order:
|
||
split_records = [
|
||
(source_row, record)
|
||
for source_row, record in zip(rows, records, strict=True)
|
||
if record["split"] == split_name
|
||
]
|
||
file_id = new_id("dfile")
|
||
version_id = new_id("dfv")
|
||
content = "".join(
|
||
json_dumps(record) + "\n" for _, record in split_records
|
||
)
|
||
raw = content.encode("utf-8")
|
||
split_specs.append(
|
||
{
|
||
"split": split_name,
|
||
"records": split_records,
|
||
"file_id": file_id,
|
||
"version_id": version_id,
|
||
"content": content,
|
||
"raw": raw,
|
||
"checksum": hashlib.sha256(raw).hexdigest(),
|
||
"storage_object_id": (
|
||
f"db://data-process/{task_id}/{file_id}/v1"
|
||
),
|
||
}
|
||
)
|
||
source_result_ids = [row["id"] for row in rows]
|
||
common_metadata = {
|
||
"source": "data_process",
|
||
"storage_backend": "database",
|
||
"source_task_id": task_id,
|
||
"output_type": _task_output_type(task),
|
||
"reasoning_detail": _task_reasoning_detail(task),
|
||
"source_file_ids": [item["id"] for item in self._source_ids(conn, task_id)],
|
||
"source_result_ids": source_result_ids,
|
||
"format": payload.get("format") or "alpaca_jsonl",
|
||
"split": requested_split,
|
||
}
|
||
|
||
existing_datasets = conn.execute(
|
||
"""
|
||
SELECT * FROM datasets
|
||
WHERE source='task' AND deleted_at IS NULL
|
||
AND (
|
||
source_task_id=%s
|
||
OR (source_task_id IS NULL AND task_id=%s)
|
||
)
|
||
ORDER BY created_at, id
|
||
""",
|
||
(task_id, task_id),
|
||
).fetchall()
|
||
existing_by_split: dict[str, dict[str, Any]] = {}
|
||
primary_existing = None
|
||
for existing in existing_datasets:
|
||
existing_metadata = _json_value(existing.get("metadata"), {})
|
||
existing_split = str(existing_metadata.get("dataset_split") or "")
|
||
if existing_split in split_order:
|
||
existing_by_split[existing_split] = existing
|
||
if str(existing["id"]) == str(task.get("output_dataset_id") or ""):
|
||
primary_existing = existing
|
||
if primary_existing and "train" not in existing_by_split:
|
||
# 兼容旧版“一个数据集包含三个文件”的发布物，原数据集复用为训练集。
|
||
existing_by_split["train"] = primary_existing
|
||
|
||
existing_group_metadata = _json_value(
|
||
(primary_existing or {}).get("metadata"), {}
|
||
)
|
||
base_dataset_name = str(
|
||
existing_group_metadata.get("base_dataset_name")
|
||
or payload["dataset_name"]
|
||
).strip()
|
||
for suffix in ("-训练集", "-验证集", "-测试集"):
|
||
if base_dataset_name.endswith(suffix):
|
||
base_dataset_name = base_dataset_name[: -len(suffix)].rstrip()
|
||
break
|
||
split_group_id = str(
|
||
existing_group_metadata.get("split_group_id")
|
||
or f"dsg_{hashlib.sha256(task_id.encode()).hexdigest()[:20]}"
|
||
)
|
||
dataset_ids = {
|
||
spec["split"]: str(existing_by_split[spec["split"]]["id"])
|
||
if spec["split"] in existing_by_split
|
||
else new_id("dataset")
|
||
for spec in split_specs
|
||
}
|
||
created_any = any(
|
||
spec["split"] not in existing_by_split for spec in split_specs
|
||
)
|
||
split_labels = {
|
||
"train": "训练集",
|
||
"validation": "验证集",
|
||
"test": "测试集",
|
||
}
|
||
dataset_types = {"train": "train", "validation": "val", "test": "test"}
|
||
published_datasets: list[dict[str, Any]] = []
|
||
try:
|
||
for spec in split_specs:
|
||
split_name = str(spec["split"])
|
||
dataset_id = dataset_ids[split_name]
|
||
existing_dataset = existing_by_split.get(split_name)
|
||
dataset_metadata = {
|
||
**common_metadata,
|
||
"base_dataset_name": base_dataset_name,
|
||
"dataset_split": split_name,
|
||
"split_group_id": split_group_id,
|
||
"split_dataset_ids": dataset_ids,
|
||
"split_counts": {
|
||
name: split_counts[name] if name == split_name else 0
|
||
for name in split_order
|
||
},
|
||
}
|
||
dataset_name = f"{base_dataset_name}-{split_labels[split_name]}"
|
||
if existing_dataset:
|
||
conn.execute(
|
||
"DELETE FROM dataset_records WHERE dataset_id=%s", (dataset_id,)
|
||
)
|
||
conn.execute(
|
||
"""DELETE FROM dataset_file_versions
|
||
WHERE dataset_file_id IN
|
||
(SELECT id FROM dataset_files WHERE dataset_id=%s)""",
|
||
(dataset_id,),
|
||
)
|
||
conn.execute(
|
||
"DELETE FROM dataset_files WHERE dataset_id=%s", (dataset_id,)
|
||
)
|
||
dataset = conn.execute(
|
||
"""
|
||
UPDATE datasets
|
||
SET name=%s, type=%s, storage_type=%s, size=%s, size_bytes=%s,
|
||
count=%s, record_count=%s, description=%s, metadata=%s,
|
||
updated_at=%s
|
||
WHERE id=%s RETURNING *
|
||
""",
|
||
(
|
||
dataset_name,
|
||
dataset_types[split_name],
|
||
payload.get("storage_type") or "local",
|
||
f"{len(spec['raw'])} B",
|
||
len(spec["raw"]),
|
||
len(spec["records"]),
|
||
len(spec["records"]),
|
||
payload.get("description") or task.get("description") or "",
|
||
json_dumps(dataset_metadata),
|
||
now,
|
||
dataset_id,
|
||
),
|
||
).fetchone()
|
||
else:
|
||
dataset = conn.execute(
|
||
"""
|
||
INSERT INTO datasets
|
||
(id, name, type, storage_type, source, task_id, source_task_id,
|
||
size, size_bytes, count, record_count, description, metadata,
|
||
tenant_id, project_id, owner_id, created_by, create_time,
|
||
created_at, updated_at)
|
||
VALUES (
|
||
%s, %s, %s, %s, 'task', %s, %s,
|
||
%s, %s, %s, %s, %s, %s,
|
||
%s, %s, %s, %s, %s, %s, %s
|
||
)
|
||
RETURNING *
|
||
""",
|
||
(
|
||
dataset_id,
|
||
dataset_name,
|
||
dataset_types[split_name],
|
||
payload.get("storage_type") or "local",
|
||
task_id,
|
||
task_id,
|
||
f"{len(spec['raw'])} B",
|
||
len(spec["raw"]),
|
||
len(spec["records"]),
|
||
len(spec["records"]),
|
||
payload.get("description") or task.get("description") or "",
|
||
json_dumps(dataset_metadata),
|
||
task.get("tenant_id"),
|
||
task.get("project_id"),
|
||
task.get("owner_id"),
|
||
payload.get("created_by") or task.get("created_by"),
|
||
now,
|
||
now,
|
||
now,
|
||
),
|
||
).fetchone()
|
||
|
||
file_metadata = {**dataset_metadata, "file_split": split_name}
|
||
version = {
|
||
"id": spec["version_id"],
|
||
"version_no": 1,
|
||
"version": 1,
|
||
"description": f"data process {split_name} publish",
|
||
"checksum_sha256": spec["checksum"],
|
||
"size_bytes": len(spec["raw"]),
|
||
"record_count": len(spec["records"]),
|
||
"created_at": now,
|
||
"create_time": now,
|
||
"source_task_id": task_id,
|
||
"storage_object_id": spec["storage_object_id"],
|
||
}
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO dataset_files
|
||
(id, dataset_id, name, storage_object_id, size, content,
|
||
active_version_id, versions, create_time, current_version_id,
|
||
size_bytes, record_count, file_format, checksum_sha256, version_no,
|
||
source_task_id, tenant_id, project_id, created_by, metadata,
|
||
created_at, updated_at)
|
||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s,
|
||
%s, %s, 1, %s, %s, %s, %s, %s, %s, %s)
|
||
""",
|
||
(
|
||
spec["file_id"],
|
||
dataset_id,
|
||
f"{base_dataset_name}.{split_name}.jsonl",
|
||
spec["storage_object_id"],
|
||
f"{len(spec['raw'])} B",
|
||
spec["content"],
|
||
spec["version_id"],
|
||
json_dumps([version]),
|
||
now,
|
||
spec["version_id"],
|
||
len(spec["raw"]),
|
||
len(spec["records"]),
|
||
"jsonl",
|
||
spec["checksum"],
|
||
task_id,
|
||
task.get("tenant_id"),
|
||
task.get("project_id"),
|
||
payload.get("created_by") or task.get("created_by"),
|
||
json_dumps(file_metadata),
|
||
now,
|
||
now,
|
||
),
|
||
)
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO dataset_file_versions
|
||
(id, dataset_file_id, version_no, storage_object_id, content_preview,
|
||
description, size_bytes, record_count, checksum_sha256,
|
||
source_task_id, metadata, created_by, created_at)
|
||
VALUES (%s, %s, 1, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
|
||
""",
|
||
(
|
||
spec["version_id"],
|
||
spec["file_id"],
|
||
spec["storage_object_id"],
|
||
spec["content"][:2000],
|
||
f"data process {split_name} publish",
|
||
len(spec["raw"]),
|
||
len(spec["records"]),
|
||
spec["checksum"],
|
||
task_id,
|
||
json_dumps(file_metadata),
|
||
payload.get("created_by") or task.get("created_by"),
|
||
now,
|
||
),
|
||
)
|
||
for line_number, (source_row, record) in enumerate(
|
||
spec["records"], start=1
|
||
):
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO dataset_records
|
||
(id, dataset_id, dataset_file_id, version_id, line_no, split,
|
||
instruction, input, output, raw, status, source_task_id,
|
||
source_result_id, preview_item_id, created_at)
|
||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s,
|
||
%s, %s, %s, %s)
|
||
""",
|
||
(
|
||
new_id("drec"),
|
||
dataset_id,
|
||
spec["file_id"],
|
||
spec["version_id"],
|
||
line_number,
|
||
record["split"],
|
||
record["instruction"],
|
||
record["input"],
|
||
record["output"],
|
||
json_dumps(
|
||
{
|
||
**record,
|
||
"source_task_id": task_id,
|
||
"source_result_id": source_row["id"],
|
||
"preview_item_id": source_row.get("preview_item_id"),
|
||
}
|
||
),
|
||
source_row["status"],
|
||
task_id,
|
||
source_row["id"],
|
||
source_row.get("preview_item_id"),
|
||
now,
|
||
),
|
||
)
|
||
published_datasets.append(_decode_row(dataset) or {})
|
||
|
||
except psycopg.errors.UniqueViolation as exc:
|
||
raise ConflictError("dataset name already exists") from exc
|
||
train_dataset_id = dataset_ids.get("train")
|
||
if not train_dataset_id:
|
||
raise InvalidStateError("published split does not contain training data")
|
||
conn.execute(
|
||
"""
|
||
UPDATE data_process_tasks
|
||
SET output_dataset_id=%s, updated_at=%s, updated_by=%s
|
||
WHERE id=%s
|
||
""",
|
||
(train_dataset_id, now, payload.get("created_by"), task_id),
|
||
)
|
||
train_dataset = next(
|
||
item
|
||
for item in published_datasets
|
||
if _json_value(item.get("metadata"), {}).get("dataset_split") == "train"
|
||
)
|
||
return {
|
||
"dataset": train_dataset,
|
||
"datasets": published_datasets,
|
||
"output_datasets": published_datasets,
|
||
"created": created_any,
|
||
"split_counts": split_counts,
|
||
}
|
||
|
||
@staticmethod
|
||
def _source_ids(
|
||
conn: psycopg.Connection[dict[str, Any]], task_id: str
|
||
) -> list[dict[str, Any]]:
|
||
return conn.execute(
|
||
"""
|
||
SELECT id FROM data_process_source_files
|
||
WHERE task_id=%s AND deleted_at IS NULL ORDER BY created_at, id
|
||
""",
|
||
(task_id,),
|
||
).fetchall()
|
||
|
||
|
||
@lru_cache
|
||
def get_data_process_store() -> DataProcessStore:
|
||
return DataProcessStore()
|