"""数据处理存储层 - 任务管理。""" from __future__ import annotations from typing import Any from collections.abc import Sequence import psycopg from .base import ( StoreBase, utcnow, new_id, repeat_task_id, json_dumps, _json_value, _decode_row, _public_task, _business_config, _preview_config_value, _preview_config_changed, _preview_config_projection, _normalized_preprocess_options, _regeneration_marker, _is_regeneration_prepared, _task_output_type, _task_reasoning_detail, _reasoning_output_is_valid, _dpo_fields_are_valid, _source_storage_descriptor, _serialize_value, NotFoundError, ConflictError, InvalidStateError, EDITABLE_STATUSES, ACTIVE_PREVIEW_STATUSES, WORKFLOW_STEPS, _REGENERATION_MARKER_KEY, _REPEAT_SOURCE_TASK_KEY, _REPEAT_REQUEST_KEY, _INTERNAL_CONFIG_KEYS, ) from ..algorithms import estimate_token_count class TasksMixin: """任务管理 Mixin。""" 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, results_confirmed, tenant_id, project_id, owner_id, created_by, updated_by, created_at, updated_at) VALUES (%s, %s, %s, 'pending', %s, %s, %s, 0, FALSE, %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(_business_config(payload.get("config"))), 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 {} @staticmethod def _repeat_response( conn: psycopg.Connection[dict[str, Any]], row: dict[str, Any], *, source_task_id: str, created: bool, ) -> dict[str, Any]: task_id = str(row["id"]) counts = conn.execute( """ SELECT (SELECT COUNT(*) FROM data_process_source_files WHERE task_id=%s AND deleted_at IS NULL) AS source_file_count, (SELECT COUNT(*) FROM data_process_preview_items WHERE task_id=%s) AS preview_count """, (task_id, task_id), ).fetchone() or {} task = _public_task(_decode_row(row)) or {} task["source_file_count"] = int(counts.get("source_file_count") or 0) task["preview_count"] = int(counts.get("preview_count") or 0) return { "task": task, "source_task_id": source_task_id, "created": created, "copied_source_file_count": task["source_file_count"], "copied_preview_count": task["preview_count"], } def find_repeated_task( self, source_task_id: str, request_id: str, ) -> dict[str, Any] | None: """查找同一幂等请求已创建的新任务。""" task_id = repeat_task_id(source_task_id, request_id) with self.connect() as conn: row = conn.execute( "SELECT * FROM data_process_tasks WHERE id=%s", (task_id,), ).fetchone() if row is None: return None decoded = _decode_row(row) or {} config = decoded.get("config") or {} if ( config.get(_REPEAT_SOURCE_TASK_KEY) != source_task_id or config.get(_REPEAT_REQUEST_KEY) != request_id ): raise ConflictError("再次生成请求与现有任务冲突") if decoded.get("deleted_at"): raise ConflictError("此次再次生成创建的任务已被删除，请重新发起") return self._repeat_response( conn, row, source_task_id=source_task_id, created=False, ) def repeat_task( self, source_task_id: str, *, expected_updated_at: str, request_id: str, file_copies: dict[str, dict[str, str]], ) -> dict[str, Any]: """复制已确认任务的配置、源文件和预览，结果与发布数据保持独立。""" task_id = repeat_task_id(source_task_id, request_id) now = utcnow() try: with self.connect() as conn: existing = conn.execute( "SELECT * FROM data_process_tasks WHERE id=%s FOR UPDATE", (task_id,), ).fetchone() if existing is not None: decoded = _decode_row(existing) or {} config = decoded.get("config") or {} if ( config.get(_REPEAT_SOURCE_TASK_KEY) != source_task_id or config.get(_REPEAT_REQUEST_KEY) != request_id ): raise ConflictError("再次生成请求与现有任务冲突") if decoded.get("deleted_at"): raise ConflictError("此次再次生成创建的任务已被删除，请重新发起") return self._repeat_response( conn, existing, source_task_id=source_task_id, created=False, ) source_task = self._task_in_connection( conn, source_task_id, for_update=True, ) if ( source_task.get("status") != "completed" or source_task.get("results_confirmed") is False ): raise InvalidStateError("只有已完成并确认结果的任务可以再次生成") if source_task.get("preview_status") in ACTIVE_PREVIEW_STATUSES: raise ConflictError("源任务仍在处理切分，暂时不能再次生成") if expected_updated_at != _serialize_value(source_task.get("updated_at")): raise ConflictError("源任务已被其他操作修改，请刷新后重试") source_files = conn.execute( """ SELECT * FROM data_process_source_files WHERE task_id=%s AND deleted_at IS NULL ORDER BY created_at, id """, (source_task_id,), ).fetchall() source_file_ids = {str(row["id"]) for row in source_files} if source_file_ids != set(file_copies): raise ConflictError("源文件快照已变化，请刷新后重试") previews = conn.execute( """ SELECT * FROM data_process_preview_items WHERE task_id=%s ORDER BY source_file_id NULLS LAST, source_start NULLS LAST, created_at, id """, (source_task_id,), ).fetchall() if not previews: raise InvalidStateError("源任务没有可用于再次生成的切分结果") suffix = f"（再次生成-{task_id[-6:]}）" base_name = str(source_task.get("name") or "数据处理任务") repeated_name = f"{base_name[: max(1, 150 - len(suffix))]}{suffix}" repeated_config = _business_config(source_task.get("config") or {}) repeated_config[_REPEAT_SOURCE_TASK_KEY] = source_task_id repeated_config[_REPEAT_REQUEST_KEY] = request_id input_count = sum(int(row.get("record_count") or 0) for row in source_files) task_row = conn.execute( """ INSERT INTO data_process_tasks (id, name, description, status, process_type, source_dataset_id, output_dataset_id, config, progress, input_count, output_count, filtered_count, duplicate_count, error_count, failure_reason, generation_run_id, results_confirmed, workflow_step, preview_status, preview_progress, preview_run_id, preview_failure_reason, preview_total_files, preview_completed_files, tenant_id, project_id, owner_id, approval_status, created_by, updated_by, created_at, updated_at) VALUES (%s, %s, %s, 'pending', %s, %s, NULL, %s, 20, %s, 0, 0, 0, 0, NULL, NULL, FALSE, 'preview', 'completed', 100, NULL, NULL, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) RETURNING * """, ( task_id, repeated_name, source_task.get("description") or "", source_task["process_type"], source_task.get("source_dataset_id"), json_dumps(repeated_config), input_count, len(source_files), len(source_files), source_task.get("tenant_id"), source_task.get("project_id"), source_task.get("owner_id"), source_task.get("approval_status") or "not_required", source_task.get("created_by"), source_task.get("created_by"), now, now, ), ).fetchone() file_id_map: dict[str, str] = {} for source in source_files: old_file_id = str(source["id"]) copy = file_copies[old_file_id] new_file_id = str(copy["id"]) storage_object_id, metadata = _source_storage_descriptor( { "storage_object_id": copy["storage_object_id"], "metadata": _json_value(source.get("metadata"), {}), }, task_id, new_file_id, ) file_id_map[old_file_id] = new_file_id 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) """, ( new_file_id, task_id, storage_object_id, source["name"], source.get("size_bytes") or 0, source.get("record_count") or 0, source.get("file_format"), source["checksum_sha256"], source.get("content") or "", source.get("content_preview"), json_dumps(metadata), source_task.get("tenant_id"), source_task.get("project_id"), source.get("created_by") or source_task.get("created_by"), now, now, ), ) for preview in previews: old_source_file_id = preview.get("source_file_id") 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) """, ( new_id("dpp"), task_id, file_id_map.get(str(old_source_file_id)) if old_source_file_id else None, preview.get("original_content") or "", preview.get("edited_content") or "", preview.get("source_start"), preview.get("source_end"), preview.get("source_start_line"), preview.get("source_end_line"), max(0, int(preview.get("token_count") or 0)), preview.get("status") or "original", json_dumps(_json_value(preview.get("quality_score"), {})), now, now, ), ) return self._repeat_response( conn, task_row or {}, source_task_id=source_task_id, created=True, ) except psycopg.errors.UniqueViolation as exc: raise ConflictError("再次生成任务名称或请求发生冲突，请重试") from exc 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("preview_status") in ACTIVE_PREVIEW_STATUSES: raise InvalidStateError("task cannot be edited while preview is running") if task.get("output_dataset_id") and not _is_regeneration_prepared(task): raise InvalidStateError("published task cannot be edited") def update_workflow_step(self, task_id: str, workflow_step: str) -> dict[str, Any]: """独立保存向导位置，不触发配置或结果失效逻辑。""" if workflow_step not in WORKFLOW_STEPS: raise ValueError("invalid data process workflow step") with self.connect() as conn: row = conn.execute( """ UPDATE data_process_tasks SET workflow_step=%s, updated_at=%s WHERE id=%s AND deleted_at IS NULL RETURNING * """, (workflow_step, utcnow(), task_id), ).fetchone() if not row: raise NotFoundError("data process task not found") return _public_task(_decode_row(row)) or {} 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 "") chosen = str(raw.get("chosen") or "") rejected = str(raw.get("rejected") or "") output = str( record.get("output") or raw.get("output") or chosen 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, chosen, rejected, original_instruction, original_input, original_output, original_chosen, original_rejected, status, error, split, quality_score, created_at, updated_at) VALUES (%s, %s, %s, %s, %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, chosen, rejected, instruction, input_text, output, chosen, rejected, 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, results_confirmed=TRUE, workflow_step='results', 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 = _business_config(payload["config"]) current_config = dict(task.get("config") or {}) for key in _INTERNAL_CONFIG_KEYS: if key in current_config: next_config[key] = current_config[key] 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, "results_confirmed": False, "preview_status": "idle", "preview_progress": 0, "preview_run_id": None, "preview_failure_reason": None, "preview_total_files": 0, "preview_completed_files": 0, "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") if task.get("preview_status") in ACTIVE_PREVIEW_STATUSES: raise ConflictError("running preview 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 = _business_config(payload.get("config")) for key in (_REPEAT_SOURCE_TASK_KEY, _REPEAT_REQUEST_KEY): if key in current_config: next_config[key] = current_config[key] 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: self._task_in_connection(conn, task_id, for_update=True) now = utcnow() conn.execute( """ UPDATE data_process_tasks SET status=CASE WHEN status='running' THEN 'stopped' ELSE status END, generation_run_id=NULL, preview_status=CASE WHEN preview_status IN ('queued', 'running') THEN 'cancelled' ELSE preview_status END, preview_run_id=NULL, deleted_at=%s, deleted_by=%s, updated_at=%s WHERE id=%s """, (now, deleted_by, now, task_id), )