ai-flow-python/app/controller/flow_controller.py

import logging
import time
from concurrent.futures import ThreadPoolExecutor
from typing import Any

from fastapi import APIRouter, Request

from app.config import settings
from app.repository.record_interface import RecordInterface
from app.service.engine.workflow_engine import WorkflowEngine
from app.service.spec.workflow_loader import load_from_directory
from app.service.task.task_executor import TaskExecutor
from app.service.template.template_service import TemplateService
from app.util.tracking import set_tracking_id

logger = logging.getLogger(__name__)

router = APIRouter()

# 全局单例（在 main.py 中初始化并注入）
_record_store: RecordInterface | None = None
_template_service: TemplateService | None = None
_engine: WorkflowEngine | None = None
_executor: ThreadPoolExecutor | None = None


def init_controller(
    record_store: RecordInterface,
    template_service: TemplateService,
):
    global _record_store, _template_service, _engine, _executor
    _record_store = record_store
    _template_service = template_service
    _engine = WorkflowEngine(TaskExecutor.create_default())
    _executor = ThreadPoolExecutor(
        max_workers=settings.executor_pool_size,
        thread_name_prefix="workflow-executor-",
    )
    logger.info("FlowController initialized with pool size: %d", settings.executor_pool_size)


def shutdown_controller():
    global _executor
    if _executor:
        _executor.shutdown(wait=True, cancel_futures=False)
        logger.info("FlowController thread pool shut down")


@router.post("/execute")
async def execute(request: Request) -> dict[str, Any]:
    req: dict = await request.json()
    template_url = req.get("template")
    inputs: dict = req.get("inputs") if isinstance(req.get("inputs"), dict) else {}

    logger.info("POST /execute  template=%s, input_keys=%s", template_url, list(inputs.keys()))

    if not template_url or not str(template_url).strip():
        logger.warning("POST /execute  rejected: template parameter is required")
        return {"state": "failed", "message": "template parameter is required"}

    task_id = _record_store.create_record(req)
    logger.info("[%s] Task accepted, submitting to thread pool", task_id)

    def _run():
        set_tracking_id(task_id)
        t0 = time.time()
        try:
            _record_store.set_processing(task_id)

            logger.info("Loading template: %s", template_url)
            try:
                template_path = _template_service.get_template_path(template_url)
                flow = load_from_directory(template_path)
                logger.info("Workflow loaded: version=%s, steps=%s, order=%s",
                            flow.version, len(flow.steps), flow.order)
            except Exception as e:
                logger.error("Failed to load workflow: %s", e)
                _record_store.set_failed(task_id, f"Failed to load workflow: {e}")
                return

            logger.info("Starting workflow execution")
            result = _engine.execute(task_id, flow, inputs)
            elapsed = time.time() - t0

            if result.success:
                logger.info("Workflow completed successfully in %.2fs", elapsed)
                _record_store.set_success(task_id, result.outputs)
            else:
                logger.warning("Workflow failed in %.2fs: %s", elapsed, result.error_message)
                _record_store.set_failed(task_id, result.error_message)
        except Exception as e:
            elapsed = time.time() - t0
            logger.error("Workflow execution exception in %.2fs: %s", elapsed, e, exc_info=True)
            _record_store.set_failed(task_id, str(e))

    _executor.submit(_run)
    return {"state": "successes", "taskId": task_id}


@router.get("/result")
async def result(taskId: str) -> dict[str, Any]:
    logger.info("GET /result  taskId=%s", taskId)
    r = _record_store.get_record(taskId)
    if r is None:
        logger.warning("GET /result  taskId=%s not found", taskId)
        return {"error": "task not found"}
    return {
        "taskId": r.task_id,
        "status": r.status,
        "message": r.message,
        "type": r.type,
        "data": r.data,
        "request": r.request or {},
        "createTime": r.create_time,
        "updateTime": r.update_time,
    }


@router.get("/tasks")
async def get_all_tasks() -> dict[str, Any]:
    all_tasks = _record_store.get_all_records()
    return {
        "total": len(all_tasks),
        "tasks": {tid: rec.to_dict() for tid, rec in all_tasks.items()},
    }


@router.post("/cleanup")
async def cleanup_tasks(request: Request) -> dict[str, Any]:
    req: dict = await request.json()
    max_age = int(req.get("maxAgeHours", 24))
    _record_store.delete_records(max_age)
    return {"message": "Cleanup completed", "maxAgeHours": max_age}


@router.get("/health")
async def health() -> dict[str, Any]:
    return {
        "status": "UP",
        "service": settings.app_name,
        "timestamp": int(time.time() * 1000),
    }


@router.get("/executor/status")
async def executor_status() -> dict[str, Any]:
    if _executor is None:
        return {"error": "executor not initialized"}
    return {
        "maxWorkers": _executor._max_workers,
        "isShutdown": _executor._shutdown,
    }


@router.get("/template/cache/status")
async def template_cache_status() -> dict[str, Any]:
    return _template_service.get_cache_stats()


@router.post("/template/cache/cleanup")
async def template_cache_cleanup(request: Request) -> dict[str, Any]:
    req: dict = await request.json()
    max_age = int(req.get("maxAgeHours", 24))
    _template_service.cleanup_expired_templates(max_age)
    return {"message": "Template cache cleanup completed", "maxAgeHours": max_age}