PMFlow/backend/tests/test_training_pipeline_service.py

from __future__ import annotations

import os

import pytest
import pytest_asyncio
from sqlalchemy import JSON
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine

os.environ["DATABASE_URL"] = "sqlite+aiosqlite:///:memory:"

from app.models.base import Base
from app.models.reply_iteration import QualityAsset, QualityAssetPair
from app.models.training_pipeline import (
    TrainedModelArtifact,
    TrainedModelDeployment,
    TrainingDataset,
    TrainingDatasetItem,
    TrainingEvaluation,
    TrainingJob,
    TrainingJobRun,
)
from app.services.training_pipeline.service import TrainingPipelineService


def _use_sqlite_json_types() -> None:
    for table in (
        QualityAsset.__table__,
        QualityAssetPair.__table__,
        TrainingDataset.__table__,
        TrainingDatasetItem.__table__,
        TrainingJob.__table__,
        TrainingJobRun.__table__,
        TrainingEvaluation.__table__,
        TrainedModelArtifact.__table__,
        TrainedModelDeployment.__table__,
    ):
        for column in table.columns:
            if column.name in {
                "source_snapshot",
                "tags",
                "selection_rules",
                "channel_scope",
                "redacted_snapshot",
                "metrics_json",
                "scope_json",
                "evaluation_summary",
                "effective_scope",
            }:
                column.type = JSON()


@pytest_asyncio.fixture
async def db_session():
    engine = create_async_engine("sqlite+aiosqlite:///:memory:")
    _use_sqlite_json_types()
    async with engine.begin() as conn:
        await conn.run_sync(
            lambda sync_conn: Base.metadata.create_all(
                sync_conn,
                tables=[
                    QualityAsset.__table__,
                    QualityAssetPair.__table__,
                    TrainingDataset.__table__,
                    TrainingDatasetItem.__table__,
                    TrainingJob.__table__,
                    TrainingJobRun.__table__,
                    TrainingEvaluation.__table__,
                    TrainedModelArtifact.__table__,
                    TrainedModelDeployment.__table__,
                ],
            )
        )
    maker = async_sessionmaker(engine, expire_on_commit=False)
    async with maker() as session:
        yield session
    await engine.dispose()


@pytest.mark.asyncio
async def test_training_dataset_rules_and_job_bootstrap(db_session) -> None:
    golden = QualityAsset(
        tenant_id=1,
        project_id=8,
        asset_type="golden",
        status="approved",
        prompt_key="PROMPT_FEISHU_DIRECT_QA",
        workflow_type="query.read_only",
        asked_via="feishu",
        question_text="现在系统走哪个层级",
        answer_text="当前走 high。",
        corrected_answer="当前走 high。",
        source_snapshot={"approved": True},
    )
    failure = QualityAsset(
        tenant_id=1,
        project_id=8,
        asset_type="failure",
        status="approved",
        prompt_key="PROMPT_FEISHU_DIRECT_QA",
        workflow_type="query.read_only",
        asked_via="feishu",
        question_text="这个回复为什么不对",
        answer_text="让我先看看。",
        corrected_answer="直接回答路由层级和原因。",
        source_snapshot={"approved": True},
    )
    db_session.add_all([golden, failure])
    await db_session.flush()

    pair = QualityAssetPair(
        tenant_id=1,
        project_id=8,
        pair_type="preference",
        status="approved",
        chosen_asset_id=golden.id,
        rejected_asset_id=failure.id,
        reason_text="黄金样本明显更符合产品要求",
    )
    db_session.add(pair)
    await db_session.commit()

    service = TrainingPipelineService(db_session)
    sft_dataset = await service.create_dataset(
        tenant_id=1,
        actor_id=9001,
        dataset_type="sft",
        version="sft-v1",
        asset_ids=[golden.id],
        pair_ids=[],
        selection_rules={"source": "manual"},
        notes="SFT 首批数据集",
    )
    preference_dataset = await service.create_dataset(
        tenant_id=1,
        actor_id=9001,
        dataset_type="preference",
        version="pref-v1",
        asset_ids=[],
        pair_ids=[pair.id],
        selection_rules={"source": "manual"},
        notes="偏好对首批数据集",
    )

    assert sft_dataset["dataset_type"] == "sft"
    assert sft_dataset["sample_count"] == 1
    assert preference_dataset["dataset_type"] == "preference"
    assert preference_dataset["sample_count"] == 1

    with pytest.raises(ValueError, match="SFT dataset only accepts approved golden assets"):
        await service.create_dataset(
            tenant_id=1,
            actor_id=9001,
            dataset_type="sft",
            version="sft-v2",
            asset_ids=[failure.id],
            pair_ids=[],
            selection_rules={"source": "manual"},
            notes="invalid",
        )

    job = await service.create_job(
        tenant_id=1,
        actor_id=9001,
        dataset_id=sft_dataset["id"],
        job_type="sft",
        base_model_family="glm",
        target_slot="high",
        provider_backend="openclaw-training",
    )
    evaluations = await service.list_evaluations(
        tenant_id=1,
        job_id=job["id"],
    )
    promoted = await service.promote_artifact(
        tenant_id=1,
        artifact_id=job["artifact"]["id"],
        actor_id=9001,
    )
    decided = await service.decide_deployment(
        tenant_id=1,
        deployment_id=promoted["id"],
        actor_id=9001,
        decision="accepted",
        notes="评测通过，允许候选部署",
    )

    assert job["status"] == "queued"
    assert job["artifact"]["target_slot"] == "high"
    assert {item["evaluation_type"] for item in evaluations["items"]} == {
        "offline_judge",
        "structure_compliance",
        "shadow_replay",
    }
    assert promoted["deployment_status"] == "candidate"
    assert decided["deployment_status"] == "active"