75 lines
2.2 KiB
Python
75 lines
2.2 KiB
Python
import json
|
||
from typing import Dict, Optional
|
||
|
||
from fastapi import APIRouter, Body, HTTPException
|
||
from fastapi.responses import JSONResponse
|
||
|
||
from workflow_service.comfy import comfy_workflow
|
||
from workflow_service.comfy.comfy_queue import queue_manager
|
||
from workflow_service.database.api import get_workflow
|
||
|
||
run_router = APIRouter(
|
||
prefix="/api/run",
|
||
tags=["Run"],
|
||
)
|
||
|
||
|
||
@run_router.post("")
|
||
async def run_workflow(
|
||
workflow_name: str,
|
||
workflow_version: Optional[str] = None,
|
||
data: Dict = Body(...),
|
||
):
|
||
"""
|
||
异步执行工作流。
|
||
立即返回任务ID，调用者可以通过任务ID查询执行状态。
|
||
"""
|
||
try:
|
||
if not workflow_name:
|
||
raise HTTPException(status_code=400, detail="`workflow_name` 字段是必需的")
|
||
|
||
# 获取工作流定义
|
||
workflow_data = await get_workflow(workflow_name, workflow_version)
|
||
if not workflow_data:
|
||
detail = (
|
||
f"工作流 '{workflow_name}'"
|
||
+ (f" 带版本 '{workflow_version}'" if workflow_version else " (最新版)")
|
||
+ " 未找到。"
|
||
)
|
||
raise HTTPException(status_code=404, detail=detail)
|
||
|
||
workflow = json.loads(workflow_data["workflow_json"])
|
||
api_spec = comfy_workflow.parse_api_spec(workflow)
|
||
|
||
# 提交到队列
|
||
workflow_run_id = await queue_manager.add_task(
|
||
workflow_name=workflow_name,
|
||
workflow_data=workflow,
|
||
api_spec=api_spec,
|
||
request_data=data,
|
||
)
|
||
|
||
return JSONResponse(
|
||
content={
|
||
"workflow_run_id": workflow_run_id,
|
||
"status": "queued",
|
||
"message": "工作流已提交到队列，正在等待执行",
|
||
},
|
||
status_code=202,
|
||
)
|
||
|
||
except Exception as e:
|
||
raise HTTPException(status_code=500, detail=f"提交工作流失败: {str(e)}")
|
||
|
||
|
||
@run_router.get("/{workflow_run_id}")
|
||
async def get_run_status(workflow_run_id: str):
|
||
"""
|
||
获取工作流执行状态。
|
||
"""
|
||
try:
|
||
status = await queue_manager.get_task_status(workflow_run_id)
|
||
return status
|
||
except Exception as e:
|
||
raise HTTPException(status_code=500, detail=f"获取状态失败: {str(e)}")
|