aiagent/backend/app/services/optimization_engine.py

"""
工作流自动优化引擎 — 根据瓶颈检测结果自动生成优化方案
"""
from __future__ import annotations

import logging
from typing import Any, Dict, List, Optional

from app.services.bottleneck_detector import bottleneck_detector

logger = logging.getLogger(__name__)


class OptimizationEngine:
    """工作流自动优化引擎 — 生成 DAG 优化方案"""

    def analyze_and_optimize(self, hours: int = 24) -> Dict[str, Any]:
        """运行完整分析并生成优化方案。"""
        analysis = bottleneck_detector.run_full_analysis(hours=hours)
        bottlenecks = analysis.get("bottlenecks", [])
        optimizations = self.generate_optimizations(bottlenecks)
        dag_changes = self.generate_dag_changes(optimizations)

        return {
            "period_hours": hours,
            "summary": analysis,
            "optimizations": optimizations,
            "dag_changes": dag_changes,
            "requires_approval": len(optimizations) > 0,
        }

    def generate_optimizations(self, bottlenecks: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """根据瓶颈生成具体优化方案。"""
        optimizations = []

        for b in bottlenecks:
            severity = b.get("severity", "low")
            if severity == "low":
                continue

            opt = {
                "node_type": b["node_type"],
                "severity": severity,
                "current_metrics": {
                    "p95_ms": b.get("p95_ms"),
                    "error_rate": b.get("error_rate"),
                    "p50_ms": b.get("p50_ms"),
                },
                "changes": [],
            }

            if b.get("is_bottleneck"):
                opt["changes"].append({
                    "type": "parallelize",
                    "description": f"将 {b['node_type']} 拆分为并行子节点",
                    "estimated_improvement": "延迟降低 50-70%",
                    "implementation": {
                        "action": "replace_node",
                        "new_structure": "parallel_gateway",
                        "sub_nodes": self._suggest_split(b["node_type"]),
                    },
                })

            if b.get("is_problematic"):
                opt["changes"].append({
                    "type": "add_validation",
                    "description": f"添加前置校验节点，失败率 {b['error_rate']:.1%}",
                    "estimated_improvement": "失败率降至 5% 以下",
                    "implementation": {
                        "action": "insert_before",
                        "new_node": {
                            "type": "validate",
                            "name": f"PreValidate_{b['node_type']}",
                            "config": {"required_fields": [], "timeout_ms": 5000},
                        },
                    },
                })
                opt["changes"].append({
                    "type": "add_retry",
                    "description": "添加重试逻辑：最多 3 次，指数退避",
                    "estimated_improvement": "间歇性失败自动恢复",
                    "implementation": {
                        "action": "update_config",
                        "config": {"retry_count": 3, "retry_delay_ms": 1000, "backoff": "exponential"},
                    },
                })

            if b.get("is_inefficient"):
                opt["changes"].append({
                    "type": "add_cache",
                    "description": f"添加结果缓存，TTL: 300s",
                    "estimated_improvement": "重复调用减少 30-60%",
                    "implementation": {
                        "action": "insert_before",
                        "new_node": {
                            "type": "cache_check",
                            "name": f"Cache_{b['node_type']}",
                            "config": {"ttl_seconds": 300, "key_fields": []},
                        },
                    },
                })

            if opt["changes"]:
                optimizations.append(opt)

        return optimizations

    def _suggest_split(self, node_type: str) -> List[Dict[str, Any]]:
        """根据节点类型建议拆分方案。"""
        suggestions = {
            "llm_call": [
                {"name": "PromptPreprocess", "type": "transform"},
                {"name": "LLMCall_Fast", "type": "llm_call", "config": {"model": "fast"}},
                {"name": "LLMCall_Accurate", "type": "llm_call", "config": {"model": "accurate"}},
                {"name": "ResultMerge", "type": "merge"},
            ],
            "api_request": [
                {"name": "RequestPrepare", "type": "transform"},
                {"name": "APICall_Core", "type": "api_request"},
                {"name": "ResponseParse", "type": "transform"},
            ],
            "data_query": [
                {"name": "QueryBuilder", "type": "transform"},
                {"name": "DBQuery", "type": "data_query"},
                {"name": "ResultFormat", "type": "transform"},
            ],
            "code_execution": [
                {"name": "SandboxSetup", "type": "setup"},
                {"name": "CodeRun", "type": "code_execution"},
                {"name": "OutputParse", "type": "transform"},
            ],
        }
        return suggestions.get(node_type, [
            {"name": "PreProcess", "type": "transform"},
            {"name": "CoreExecute", "type": "execute"},
            {"name": "PostProcess", "type": "transform"},
        ])

    def generate_dag_changes(self, optimizations: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """将优化方案转换为 DAG 变更列表（可用于 diff 展示）。"""
        changes = []
        for opt in optimizations:
            for change in opt.get("changes", []):
                impl = change.get("implementation", {})
                changes.append({
                    "node_type": opt["node_type"],
                    "severity": opt["severity"],
                    "action": impl.get("action"),
                    "description": change["description"],
                    "estimated_improvement": change.get("estimated_improvement"),
                    "detail": impl,
                })
        return changes

    def apply_optimization(self, workflow_id: str, optimization_id: str,
                           approved_changes: List[str]) -> Dict[str, Any]:
        """应用用户确认的优化变更（创建新版本）。"""
        # 这里与 workflow_version 系统集成，创建优化后的新版本
        # 当前返回框架结果，实际集成在后续完成
        return {
            "workflow_id": workflow_id,
            "status": "pending",
            "message": f"优化方案已记录，{len(approved_changes)} 项变更待应用",
            "applied_changes": approved_changes,
        }


optimization_engine = OptimizationEngine()