from __future__ import annotations import hashlib import json import uuid from contextlib import contextmanager from datetime import date, datetime, timezone from functools import lru_cache from pathlib import Path from typing import Any, Iterator, Sequence 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 TASK_STATUSES = {"pending", "running", "completed", "failed", "stopped"} EDITABLE_STATUSES = {"pending", "failed", "stopped", "completed"} class DataProcessStoreError(RuntimeError): pass class NotFoundError(DataProcessStoreError): pass class ConflictError(DataProcessStoreError): pass class InvalidStateError(DataProcessStoreError): pass def utcnow() -> str: return datetime.now(timezone.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 _serialize_value(value: Any) -> Any: if isinstance(value, (datetime, date)): return value.isoformat().replace("+00:00", "Z") return value 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": [], }.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: with 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 = ["deleted_at IS NULL"] params: list[Any] = [] if keyword: clauses.append("(name ILIKE %s OR COALESCE(description, '') ILIKE %s)") pattern = f"%{keyword.strip()}%" params.extend([pattern, pattern]) if status: clauses.append("status = %s") params.append(status) if process_type: clauses.append("process_type = %s") params.append(process_type) if tenant_id: clauses.append("tenant_id = %s") params.append(tenant_id) if project_id: clauses.append("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 WHERE {where}", params ).fetchone()["count"] rows = conn.execute( f""" SELECT * FROM data_process_tasks WHERE {where} ORDER BY created_at DESC, id DESC 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 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(payload.get("config") or {}), 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 _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: 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 {} 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"): raise InvalidStateError("published task cannot be edited") 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 payload.get("config") is not None: values["config"] = json_dumps(payload["config"]) if not values: 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) 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: 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 _decode_row(row) or {} 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, ) -> 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, } ], )[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 = new_id("dpsf") storage_object_id = f"db://data-process/{task_id}/{file_id}/v1" metadata_payload = { "storage_backend": "database", **(payload.get("metadata") or {}), } 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 {}) 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,)) 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, 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 """, (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_results WHERE task_id=%s", (task_id,)) conn.execute( "DELETE FROM data_process_preview_items WHERE source_file_id=%s", (file_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, utcnow(), task_id), ) def replace_preview_items( self, task_id: str, items: Sequence[dict[str, Any]] ) -> list[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) 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,)) 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 {}) 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_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_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_id, utcnow()) def _invalidate_results( self, conn: psycopg.Connection[dict[str, Any]], task_id: str, now: str, ) -> None: 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) if task.get("output_dataset_id"): 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") row = conn.execute( """ UPDATE data_process_tasks SET status='running', progress=30, failure_reason=NULL, started_at=%s, completed_at=NULL, filtered_count=0, duplicate_count=0, error_count=0, generation_run_id=%s, updated_at=%s WHERE id=%s RETURNING * """, (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 {} hard_valid = bool( str(merged.get("instruction") or "").strip() and str(merged.get("output") or "").strip() ) 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 values["error"] = ", ".join(str(flag) for flag in flags or []) 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 task.get("output_dataset_id"): dataset = conn.execute( "SELECT * FROM datasets WHERE id=%s", (task["output_dataset_id"],) ).fetchone() if dataset: return {"dataset": _decode_row(dataset), "created": False} # 数据集被外部流程清理后，解除断链并重新发布。 conn.execute( "UPDATE data_process_tasks SET output_dataset_id=NULL WHERE id=%s", (task_id,), ) 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() ) if invalid_count: raise InvalidStateError(f"task contains {invalid_count} invalid results") dataset_id = new_id("dataset") file_id = new_id("dfile") version_id = new_id("dfv") now = utcnow() requested_split = payload.get("split") or { "train": 80, "validation": 10, "test": 10, } records = [ { "instruction": row["instruction"], "input": row["input"], "output": row["output"], "split": stable_split( str(row["id"]), requested_split, seed=task_id, ), } for row in rows ] content = "".join(json_dumps(record) + "\n" for record in records) raw = content.encode("utf-8") 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] metadata = { "source": "data_process", "storage_backend": "database", "storage_object_id": storage_object_id, "source_task_id": task_id, "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": payload.get("split") or {}, } try: 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, payload["dataset_name"], payload.get("dataset_type") or "train", payload.get("storage_type") or "local", task_id, task_id, f"{len(raw)} B", len(raw), len(records), len(records), payload.get("description") or task.get("description") or "", json_dumps(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() version = { "id": version_id, "version_no": 1, "description": "data process publish", "checksum_sha256": checksum, "size_bytes": len(raw), "record_count": len(records), "created_at": now, "source_task_id": task_id, "storage_object_id": 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) """, ( file_id, dataset_id, f"{payload['dataset_name']}.jsonl", storage_object_id, f"{len(raw)} B", content, version_id, json_dumps([version]), now, version_id, len(raw), len(records), "jsonl", checksum, task_id, task.get("tenant_id"), task.get("project_id"), payload.get("created_by") or task.get("created_by"), json_dumps(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) """, ( version_id, file_id, storage_object_id, content[:2000], "data process publish", len(raw), len(records), checksum, task_id, json_dumps(metadata), payload.get("created_by") or task.get("created_by"), now, ), ) for line_number, (source_row, record) in enumerate( zip(rows, records, strict=True), 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, file_id, 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, ), ) except psycopg.errors.UniqueViolation as exc: raise ConflictError("dataset name already exists") from exc conn.execute( """ UPDATE data_process_tasks SET output_dataset_id=%s, updated_at=%s, updated_by=%s WHERE id=%s """, (dataset_id, now, payload.get("created_by"), task_id), ) return {"dataset": _decode_row(dataset), "created": True} @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()