"""数据处理存储层 - 基础设施。 包含:异常类、工具函数、常量定义、基类。 """ from __future__ import annotations import hashlib import json import re import uuid from collections.abc import Iterator from contextlib import contextmanager from datetime import UTC, date, datetime from decimal import Decimal from pathlib import Path from typing import Any import psycopg from psycopg.rows import dict_row from app.core.config import get_settings # 常量定义 TASK_STATUSES = {"pending", "running", "completed", "failed", "stopped"} EDITABLE_STATUSES = {"pending", "failed", "stopped", "completed"} ACTIVE_PREVIEW_STATUSES = {"queued", "running"} WORKFLOW_STEPS = {"create", "model", "upload", "preview", "generate", "results"} _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" _REPEAT_SOURCE_TASK_KEY = "_repeat_source_task_id" _REPEAT_REQUEST_KEY = "_repeat_request_id" _INTERNAL_CONFIG_KEYS = { _REGENERATION_MARKER_KEY, _REPEAT_SOURCE_TASK_KEY, _REPEAT_REQUEST_KEY, } # 异常类 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 repeat_task_id(source_task_id: str, request_id: str) -> str: """按源任务和请求幂等键生成稳定的新任务 ID。""" digest = hashlib.sha256(f"{source_task_id}:{request_id}".encode()).hexdigest() return f"dpt_{digest[: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*\s*(?P[\s\S]*?)\s*\s*(?P[\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 (" bool: chosen = str(row.get("chosen") or "").strip() rejected = str(row.get("rejected") or "").strip() return bool(chosen and rejected and chosen != rejected) 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 _business_config(config: dict[str, Any] | None) -> dict[str, Any]: """过滤只供服务端维护的工作流标记。""" return { key: value for key, value in (config or {}).items() if key not in _INTERNAL_CONFIG_KEYS } 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): public["config"] = _business_config(config) 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 StoreBase: """数据处理存储基类。""" 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[3] / "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)