"""数据处理存储层 - 生成管理。""" 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, DataProcessStoreError, NotFoundError, ConflictError, InvalidStateError, EDITABLE_STATUSES, ACTIVE_PREVIEW_STATUSES, WORKFLOW_STEPS, _REGENERATION_MARKER_KEY, _REPEAT_SOURCE_TASK_KEY, _REPEAT_REQUEST_KEY, ) class GenerationMixin: """生成管理 Mixin。""" 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, results_confirmed=FALSE, 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") if task.get("preview_status") in ACTIVE_PREVIEW_STATUSES: raise ConflictError("preview is still 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, results_confirmed=FALSE, workflow_step='generate', 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: with self.connect() as conn: row = conn.execute( """ SELECT status, generation_run_id FROM data_process_tasks WHERE id=%s AND deleted_at IS NULL """, (task_id,), ).fetchone() return bool( row and row.get("status") == "running" and row.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, 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, %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("chosen") or "", result.get("rejected") 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("original_chosen", result.get("chosen") or ""), result.get("original_rejected", result.get("rejected") 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, results_confirmed=FALSE, 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"), "results_confirmed": bool(task.get("results_confirmed")), "started_at": task.get("started_at"), "completed_at": task.get("completed_at"), }