YG_FT/backend/app/modules/data_process/store/base.py

"""数据处理存储层 - 基础设施。

包含：异常类、工具函数、常量定义、基类。
"""

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*<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 _dpo_fields_are_valid(row: dict[str, Any]) -> 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_minio_prefix = f"minio://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.startswith(expected_minio_prefix) and len(storage_object_id) > len(
        expected_minio_prefix
    ):
        storage_backend = "minio"
    elif storage_object_id == expected_database_reference:
        storage_backend = "database"
    elif storage_object_id.startswith(("local://data-process/", "minio://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)