PMFlow/backend/app/api/ai_steward_api/routes_training_pipeline.py

from __future__ import annotations

from fastapi import Depends, HTTPException, Query
from sqlalchemy.ext.asyncio import AsyncSession

from ...core.response import success_response
from ...dependencies.permissions import require_manager_user
from ...dependencies.tenant import get_tenant_id
from ...models import User
from ...schemas.training_pipeline import (
    TrainingDatasetCreateRequest,
    TrainingDeploymentDecisionRequest,
    TrainingJobCreateRequest,
)
from ...services.training_pipeline import TrainingPipelineService
from .shared import *


@router.post("/training/datasets")
async def create_training_dataset(
    request: TrainingDatasetCreateRequest,
    current_user: User = Depends(require_manager_user),
    db: AsyncSession = Depends(get_db),
    tenant_id: int = Depends(get_tenant_id),
):
    effective_tenant_id = int(getattr(current_user, "tenant_id", tenant_id) or tenant_id)
    try:
        data = await TrainingPipelineService(db).create_dataset(
            tenant_id=effective_tenant_id,
            actor_id=int(current_user.id),
            dataset_type=request.dataset_type,
            version=request.version,
            asset_ids=request.asset_ids,
            pair_ids=request.pair_ids,
            selection_rules=request.selection_rules,
            notes=request.notes,
        )
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    return success_response(data)


@router.get("/training/datasets")
async def list_training_datasets(
    dataset_type: str | None = Query(default=None),
    limit: int = Query(default=20, ge=1, le=100),
    offset: int = Query(default=0, ge=0),
    current_user: User = Depends(require_manager_user),
    db: AsyncSession = Depends(get_db),
    tenant_id: int = Depends(get_tenant_id),
):
    effective_tenant_id = int(getattr(current_user, "tenant_id", tenant_id) or tenant_id)
    return success_response(
        await TrainingPipelineService(db).list_datasets(
            tenant_id=effective_tenant_id,
            dataset_type=dataset_type,
            limit=limit,
            offset=offset,
        )
    )


@router.post("/training/jobs")
async def create_training_job(
    request: TrainingJobCreateRequest,
    current_user: User = Depends(require_manager_user),
    db: AsyncSession = Depends(get_db),
    tenant_id: int = Depends(get_tenant_id),
):
    effective_tenant_id = int(getattr(current_user, "tenant_id", tenant_id) or tenant_id)
    try:
        data = await TrainingPipelineService(db).create_job(
            tenant_id=effective_tenant_id,
            actor_id=int(current_user.id),
            dataset_id=request.dataset_id,
            job_type=request.job_type,
            base_model_family=request.base_model_family,
            target_slot=request.target_slot,
            provider_backend=request.provider_backend,
        )
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    return success_response(data)


@router.get("/training/jobs")
async def list_training_jobs(
    status: str | None = Query(default=None),
    limit: int = Query(default=20, ge=1, le=100),
    offset: int = Query(default=0, ge=0),
    current_user: User = Depends(require_manager_user),
    db: AsyncSession = Depends(get_db),
    tenant_id: int = Depends(get_tenant_id),
):
    effective_tenant_id = int(getattr(current_user, "tenant_id", tenant_id) or tenant_id)
    return success_response(
        await TrainingPipelineService(db).list_jobs(
            tenant_id=effective_tenant_id,
            status=status,
            limit=limit,
            offset=offset,
        )
    )


@router.get("/training/jobs/{job_id}/evaluations")
async def list_training_job_evaluations(
    job_id: int,
    current_user: User = Depends(require_manager_user),
    db: AsyncSession = Depends(get_db),
    tenant_id: int = Depends(get_tenant_id),
):
    effective_tenant_id = int(getattr(current_user, "tenant_id", tenant_id) or tenant_id)
    return success_response(
        await TrainingPipelineService(db).list_evaluations(
            tenant_id=effective_tenant_id,
            job_id=job_id,
        )
    )


@router.post("/training/artifacts/{artifact_id}/promote")
async def promote_trained_model_artifact(
    artifact_id: int,
    current_user: User = Depends(require_manager_user),
    db: AsyncSession = Depends(get_db),
    tenant_id: int = Depends(get_tenant_id),
):
    effective_tenant_id = int(getattr(current_user, "tenant_id", tenant_id) or tenant_id)
    try:
        data = await TrainingPipelineService(db).promote_artifact(
            tenant_id=effective_tenant_id,
            artifact_id=artifact_id,
            actor_id=int(current_user.id),
        )
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    return success_response(data)


@router.post("/training/deployments/{deployment_id}/decision")
async def decide_trained_model_deployment(
    deployment_id: int,
    request: TrainingDeploymentDecisionRequest,
    current_user: User = Depends(require_manager_user),
    db: AsyncSession = Depends(get_db),
    tenant_id: int = Depends(get_tenant_id),
):
    effective_tenant_id = int(getattr(current_user, "tenant_id", tenant_id) or tenant_id)
    try:
        data = await TrainingPipelineService(db).decide_deployment(
            tenant_id=effective_tenant_id,
            deployment_id=deployment_id,
            actor_id=int(current_user.id),
            decision=request.decision,
            notes=request.notes,
        )
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    return success_response(data)


@router.get("/training/deployments")
async def list_training_deployments(
    deployment_status: str | None = Query(default=None),
    limit: int = Query(default=20, ge=1, le=100),
    offset: int = Query(default=0, ge=0),
    current_user: User = Depends(require_manager_user),
    db: AsyncSession = Depends(get_db),
    tenant_id: int = Depends(get_tenant_id),
):
    effective_tenant_id = int(getattr(current_user, "tenant_id", tenant_id) or tenant_id)
    return success_response(
        await TrainingPipelineService(db).list_deployments(
            tenant_id=effective_tenant_id,
            deployment_status=deployment_status,
            limit=limit,
            offset=offset,
        )
    )