cluster4npu/core/templates/manager.py

"""
core/templates/manager.py

TemplateManager — 提供常見使用情境的預設 Pipeline 範本。

設計重點：
- 三個內建範本（yolov5_detection、fire_detection、dual_model_cascade）以常數定義。
- save_as_template 將自訂範本儲存於記憶體（in-memory），不持久化到磁碟。
- load_template 先查內建範本，再查自訂範本；找不到時拋出 ValueError。
- nodes/connections 格式與 .mflow JSON 相同（id、type 為必要欄位）。
"""
from __future__ import annotations

import time
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional


@dataclass
class PipelineTemplate:
    """單一 Pipeline 範本。

    屬性：
        template_id: 唯一識別碼（內建範本使用語意名稱；自訂範本以 custom_ 開頭）。
        name: 顯示名稱（如 "YOLOv5 物件偵測"）。
        description: 範本說明。
        nodes: 節點定義列表，格式與 .mflow 相同，每個節點至少含 id 和 type。
        connections: 連線定義列表，每條連線含 from 和 to。
    """
    template_id: str
    name: str
    description: str
    nodes: List[Dict[str, Any]]
    connections: List[Dict[str, Any]]


# ---------------------------------------------------------------------------
# 內建範本定義
# ---------------------------------------------------------------------------

_BUILTIN_TEMPLATES: List[PipelineTemplate] = [
    PipelineTemplate(
        template_id="yolov5_detection",
        name="YOLOv5 物件偵測",
        description="標準 YOLOv5 物件偵測流程：輸入影像經前處理後送入模型，後處理輸出邊界框結果。",
        nodes=[
            {"id": "input_0",      "type": "Input",       "label": "Input"},
            {"id": "preprocess_0", "type": "Preprocess",  "label": "Preprocess"},
            {"id": "model_0",      "type": "Model",       "label": "Model"},
            {"id": "postprocess_0","type": "Postprocess", "label": "Postprocess"},
            {"id": "output_0",     "type": "Output",      "label": "Output"},
        ],
        connections=[
            {"from": "input_0",       "to": "preprocess_0"},
            {"from": "preprocess_0",  "to": "model_0"},
            {"from": "model_0",       "to": "postprocess_0"},
            {"from": "postprocess_0", "to": "output_0"},
        ],
    ),
    PipelineTemplate(
        template_id="fire_detection",
        name="火焰偵測分類",
        description="火焰偵測流程：影像直接送入模型推論，後處理輸出火焰偵測結果（無前處理節點）。",
        nodes=[
            {"id": "input_0",      "type": "Input",       "label": "Input"},
            {"id": "model_0",      "type": "Model",       "label": "Model"},
            {"id": "postprocess_0","type": "Postprocess", "label": "Postprocess"},
            {"id": "output_0",     "type": "Output",      "label": "Output"},
        ],
        connections=[
            {"from": "input_0",       "to": "model_0"},
            {"from": "model_0",       "to": "postprocess_0"},
            {"from": "postprocess_0", "to": "output_0"},
        ],
    ),
    PipelineTemplate(
        template_id="dual_model_cascade",
        name="雙模型串接",
        description=(
            "兩個模型串接的複合推論流程：第一個模型的輸出結果經後處理後，"
            "作為第二個模型的輸入，適合先偵測後分類的使用情境。"
        ),
        nodes=[
            {"id": "input_0",       "type": "Input",       "label": "Input"},
            {"id": "model_0",       "type": "Model",       "label": "Model 1"},
            {"id": "postprocess_0", "type": "Postprocess", "label": "Postprocess 1"},
            {"id": "model_1",       "type": "Model",       "label": "Model 2"},
            {"id": "postprocess_1", "type": "Postprocess", "label": "Postprocess 2"},
            {"id": "output_0",      "type": "Output",      "label": "Output"},
        ],
        connections=[
            {"from": "input_0",       "to": "model_0"},
            {"from": "model_0",       "to": "postprocess_0"},
            {"from": "postprocess_0", "to": "model_1"},
            {"from": "model_1",       "to": "postprocess_1"},
            {"from": "postprocess_1", "to": "output_0"},
        ],
    ),
]

# 以 template_id 建立快速查找字典
_BUILTIN_BY_ID: Dict[str, PipelineTemplate] = {
    t.template_id: t for t in _BUILTIN_TEMPLATES
}


# ---------------------------------------------------------------------------
# TemplateManager
# ---------------------------------------------------------------------------

class TemplateManager:
    """管理內建與自訂 Pipeline 範本。

    自訂範本儲存於記憶體，每個 TemplateManager 實例各自獨立。
    """

    def __init__(self) -> None:
        # 自訂範本字典：{template_id: PipelineTemplate}
        self._custom: Dict[str, PipelineTemplate] = {}

    # ------------------------------------------------------------------
    # 公開介面
    # ------------------------------------------------------------------

    def get_builtin_templates(self) -> List[PipelineTemplate]:
        """回傳所有內建範本的清單（共 3 個）。

        回傳：
            PipelineTemplate 列表（不含自訂範本）。
        """
        return list(_BUILTIN_TEMPLATES)

    def load_template(self, template_id: str) -> PipelineTemplate:
        """依 template_id 載入範本。

        查找順序：內建範本 → 自訂範本。

        參數：
            template_id: 範本唯一識別碼。

        回傳：
            對應的 PipelineTemplate。

        引發：
            ValueError: 當 template_id 不存在於任何範本時。
        """
        if template_id in _BUILTIN_BY_ID:
            return _BUILTIN_BY_ID[template_id]

        if template_id in self._custom:
            return self._custom[template_id]

        raise ValueError(f"Template {template_id} not found")

    def save_as_template(
        self,
        pipeline_config: Dict[str, Any],
        name: str,
        description: str,
    ) -> PipelineTemplate:
        """將 Pipeline 設定儲存為新的自訂範本。

        參數：
            pipeline_config: 包含 nodes 和 connections 列表的字典。
            name: 範本顯示名稱。
            description: 範本說明。

        回傳：
            新建立的 PipelineTemplate（template_id 以 custom_ 開頭）。
        """
        safe_name = name.lower().replace(" ", "_")
        template_id = f"custom_{safe_name}_{int(time.time() * 1000)}"

        template = PipelineTemplate(
            template_id=template_id,
            name=name,
            description=description,
            nodes=list(pipeline_config.get("nodes", [])),
            connections=list(pipeline_config.get("connections", [])),
        )
        self._custom[template_id] = template
        return template