PMFlow/backend/app/api/pmflow_launch_api/__init__.py

125 lines
5.2 KiB
Python

from __future__ import annotations
import asyncio
from fastapi import APIRouter, Depends
from fastapi.responses import JSONResponse
from sqlalchemy.ext.asyncio import AsyncSession
from ...core.database import AsyncSessionLocal, get_db
from ...core.response import error_response, success_response
from ...dependencies import get_current_user
from ...models import User
from ...schemas.pmflow_launch import PMFlowLaunchConfirmRequest, PMFlowLaunchDraftJobRequest, PMFlowLaunchDraftRequest
from ...services.pmflow_launch_service import PMFlowLaunchError, PMFlowLaunchService
router = APIRouter(prefix="/pmflow-launch", tags=["pmflow-launch"])
async def _run_launch_draft_job(*, tenant_id: int, job_id: str) -> None:
async with AsyncSessionLocal() as bg_db:
service = PMFlowLaunchService(bg_db, tenant_id=int(tenant_id))
await service.run_draft_job(job_id=job_id)
@router.post("/drafts")
async def create_pmflow_launch_draft(
body: PMFlowLaunchDraftRequest,
db: AsyncSession = Depends(get_db),
current_user: User = Depends(get_current_user),
):
service = PMFlowLaunchService(db, tenant_id=int(current_user.tenant_id))
try:
payload = await service.create_draft(body=body, user_id=int(current_user.id))
except PMFlowLaunchError as exc:
return JSONResponse(status_code=exc.status_code, content=error_response(exc.status_code, f"{exc.code}: {exc.message}"))
return success_response(payload.model_dump(mode="json"))
@router.post("/draft-jobs")
async def create_pmflow_launch_draft_job(
body: PMFlowLaunchDraftJobRequest,
db: AsyncSession = Depends(get_db),
current_user: User = Depends(get_current_user),
):
service = PMFlowLaunchService(db, tenant_id=int(current_user.tenant_id))
try:
payload = await service.create_draft_job(body=body, user_id=int(current_user.id))
except PMFlowLaunchError as exc:
return JSONResponse(status_code=exc.status_code, content=error_response(exc.status_code, f"{exc.code}: {exc.message}"))
asyncio.create_task(
_run_launch_draft_job(tenant_id=int(current_user.tenant_id), job_id=payload.job_id),
name=f"pmflow-launch-draft-job-{payload.job_id}",
)
return success_response(payload.model_dump(mode="json"))
@router.get("/draft-jobs/{job_id}")
async def get_pmflow_launch_draft_job(
job_id: str,
db: AsyncSession = Depends(get_db),
current_user: User = Depends(get_current_user),
):
service = PMFlowLaunchService(db, tenant_id=int(current_user.tenant_id))
try:
payload = await service.get_draft_job(job_id=job_id, user_id=int(current_user.id))
except PMFlowLaunchError as exc:
return JSONResponse(status_code=exc.status_code, content=error_response(exc.status_code, f"{exc.code}: {exc.message}"))
return success_response(payload.model_dump(mode="json"))
@router.post("/draft-jobs/{job_id}/cancel")
async def cancel_pmflow_launch_draft_job(
job_id: str,
db: AsyncSession = Depends(get_db),
current_user: User = Depends(get_current_user),
):
service = PMFlowLaunchService(db, tenant_id=int(current_user.tenant_id))
try:
payload = await service.cancel_draft_job(job_id=job_id, user_id=int(current_user.id))
except PMFlowLaunchError as exc:
return JSONResponse(status_code=exc.status_code, content=error_response(exc.status_code, f"{exc.code}: {exc.message}"))
return success_response(payload.model_dump(mode="json"))
@router.post("/drafts/{draft_id}/revise")
async def revise_pmflow_launch_draft(
draft_id: str,
body: PMFlowLaunchDraftJobRequest,
db: AsyncSession = Depends(get_db),
current_user: User = Depends(get_current_user),
):
mode = body.mode if body.mode in {"merge_supplement", "regenerate"} else "merge_supplement"
normalized_body = body.model_copy(update={"parent_draft_id": draft_id, "mode": mode})
service = PMFlowLaunchService(db, tenant_id=int(current_user.tenant_id))
try:
payload = await service.create_draft_job(body=normalized_body, user_id=int(current_user.id))
except PMFlowLaunchError as exc:
return JSONResponse(status_code=exc.status_code, content=error_response(exc.status_code, f"{exc.code}: {exc.message}"))
asyncio.create_task(
_run_launch_draft_job(tenant_id=int(current_user.tenant_id), job_id=payload.job_id),
name=f"pmflow-launch-draft-job-{payload.job_id}",
)
return success_response(payload.model_dump(mode="json"))
@router.post("/drafts/{draft_id}/confirm")
async def confirm_pmflow_launch_draft(
draft_id: str,
body: PMFlowLaunchConfirmRequest,
db: AsyncSession = Depends(get_db),
current_user: User = Depends(get_current_user),
):
if not body.confirm:
return JSONResponse(status_code=400, content=error_response(400, "confirm_required: confirm must be true"))
service = PMFlowLaunchService(db, tenant_id=int(current_user.tenant_id))
try:
payload = await service.confirm_draft(
draft_id=draft_id,
edited_draft=body.edited_draft,
user_id=int(current_user.id),
)
except PMFlowLaunchError as exc:
return JSONResponse(status_code=exc.status_code, content=error_response(exc.status_code, f"{exc.code}: {exc.message}"))
return success_response(payload.model_dump(mode="json"))