PMFlow/backend/app/models/training_pipeline.py

"""训练 / 微调管线相关模型。"""

from __future__ import annotations

from typing import Any, Optional

from sqlalchemy import BigInteger, ForeignKey, Integer, Text, UniqueConstraint, VARCHAR
from sqlalchemy.orm import Mapped, mapped_column

from .base import Base
from .types import JSONB


class TrainingDataset(Base):
    __tablename__ = "training_datasets"

    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    project_id: Mapped[Optional[int]] = mapped_column(BigInteger, nullable=True, index=True)
    dataset_type: Mapped[str] = mapped_column(VARCHAR(16), nullable=False, index=True)
    version: Mapped[str] = mapped_column(VARCHAR(64), nullable=False, index=True)
    status: Mapped[str] = mapped_column(VARCHAR(16), nullable=False, default="ready", index=True)
    sample_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
    selection_rules: Mapped[Optional[Any]] = mapped_column(JSONB, nullable=True)
    channel_scope: Mapped[Optional[Any]] = mapped_column(JSONB, nullable=True)
    notes: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    created_by: Mapped[Optional[int]] = mapped_column(BigInteger, nullable=True, index=True)

    __table_args__ = (
        UniqueConstraint("tenant_id", "dataset_type", "version", name="uq_training_datasets_tenant_type_version"),
    )


class TrainingDatasetItem(Base):
    __tablename__ = "training_dataset_items"

    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    dataset_id: Mapped[int] = mapped_column(
        Integer,
        ForeignKey("training_datasets.id", ondelete="CASCADE"),
        nullable=False,
        index=True,
    )
    quality_asset_id: Mapped[Optional[int]] = mapped_column(
        Integer,
        ForeignKey("quality_assets.id", ondelete="CASCADE"),
        nullable=True,
        index=True,
    )
    quality_asset_pair_id: Mapped[Optional[int]] = mapped_column(
        Integer,
        ForeignKey("quality_asset_pairs.id", ondelete="CASCADE"),
        nullable=True,
        index=True,
    )
    redacted_snapshot: Mapped[Optional[Any]] = mapped_column(JSONB, nullable=True)


class TrainingJob(Base):
    __tablename__ = "training_jobs"

    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    project_id: Mapped[Optional[int]] = mapped_column(BigInteger, nullable=True, index=True)
    dataset_id: Mapped[int] = mapped_column(
        Integer,
        ForeignKey("training_datasets.id", ondelete="CASCADE"),
        nullable=False,
        index=True,
    )
    job_type: Mapped[str] = mapped_column(VARCHAR(16), nullable=False, index=True)
    base_model_family: Mapped[str] = mapped_column(VARCHAR(64), nullable=False, index=True)
    target_slot: Mapped[str] = mapped_column(VARCHAR(16), nullable=False, index=True)
    status: Mapped[str] = mapped_column(VARCHAR(16), nullable=False, default="queued", index=True)
    provider_backend: Mapped[str] = mapped_column(VARCHAR(64), nullable=False)
    artifact_ref: Mapped[Optional[str]] = mapped_column(VARCHAR(255), nullable=True)
    metrics_json: Mapped[Optional[Any]] = mapped_column(JSONB, nullable=True)
    created_by: Mapped[Optional[int]] = mapped_column(BigInteger, nullable=True, index=True)


class TrainingJobRun(Base):
    __tablename__ = "training_job_runs"

    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    training_job_id: Mapped[int] = mapped_column(
        Integer,
        ForeignKey("training_jobs.id", ondelete="CASCADE"),
        nullable=False,
        index=True,
    )
    status: Mapped[str] = mapped_column(VARCHAR(16), nullable=False, default="queued", index=True)
    provider_run_ref: Mapped[Optional[str]] = mapped_column(VARCHAR(255), nullable=True)
    metrics_json: Mapped[Optional[Any]] = mapped_column(JSONB, nullable=True)


class TrainingEvaluation(Base):
    __tablename__ = "training_evaluations"

    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    training_job_id: Mapped[int] = mapped_column(
        Integer,
        ForeignKey("training_jobs.id", ondelete="CASCADE"),
        nullable=False,
        index=True,
    )
    evaluation_type: Mapped[str] = mapped_column(VARCHAR(32), nullable=False, index=True)
    status: Mapped[str] = mapped_column(VARCHAR(16), nullable=False, default="pending", index=True)
    score: Mapped[Optional[float]] = mapped_column(nullable=True)
    evaluation_summary: Mapped[Optional[Any]] = mapped_column(JSONB, nullable=True)


class TrainedModelArtifact(Base):
    __tablename__ = "trained_model_artifacts"

    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    training_job_id: Mapped[int] = mapped_column(
        Integer,
        ForeignKey("training_jobs.id", ondelete="CASCADE"),
        nullable=False,
        index=True,
    )
    artifact_ref: Mapped[str] = mapped_column(VARCHAR(255), nullable=False, index=True)
    base_model_family: Mapped[str] = mapped_column(VARCHAR(64), nullable=False, index=True)
    target_slot: Mapped[str] = mapped_column(VARCHAR(16), nullable=False, index=True)
    status: Mapped[str] = mapped_column(VARCHAR(16), nullable=False, default="pending_eval", index=True)
    metrics_json: Mapped[Optional[Any]] = mapped_column(JSONB, nullable=True)


class TrainedModelDeployment(Base):
    __tablename__ = "trained_model_deployments"

    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    artifact_id: Mapped[int] = mapped_column(
        Integer,
        ForeignKey("trained_model_artifacts.id", ondelete="CASCADE"),
        nullable=False,
        index=True,
    )
    deployment_status: Mapped[str] = mapped_column(VARCHAR(16), nullable=False, default="candidate", index=True)
    bound_tier: Mapped[str] = mapped_column(VARCHAR(16), nullable=False, index=True)
    effective_scope: Mapped[Optional[Any]] = mapped_column(JSONB, nullable=True)
    decision: Mapped[Optional[str]] = mapped_column(VARCHAR(24), nullable=True)
    notes: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    decided_by: Mapped[Optional[int]] = mapped_column(BigInteger, nullable=True, index=True)
    decided_at: Mapped[Optional[str]] = mapped_column(VARCHAR(64), nullable=True)


__all__ = [
    "TrainingDataset",
    "TrainingDatasetItem",
    "TrainingJob",
    "TrainingJobRun",
    "TrainingEvaluation",
    "TrainedModelArtifact",
    "TrainedModelDeployment",
]