用多 Agent 系统接管工程支持：Grab 数据仓库平台的实践

import json
import time
from dataclasses import dataclass, field
from typing import Optional

# ── 1. 定义结构化上下文（Agent 之间传递信息的契约）──

@dataclass
class InvestigationResult:
    root_cause: str
    evidence: list[str] = field(default_factory=list)
    confidence: float = 0.0  # 0-1
    needs_human: bool = False

@dataclass
class EnhancementPlan:
    actions: list[str] = field(default_factory=list)
    generated_script: str = ""
    doc_update: str = ""

# ── 2. Agent 基类 ──

class Agent:
    def __init__(self, name: str):
        self.name = name

    def run(self, task: str, context: dict) -> dict:
        """子类实现具体逻辑，这里用 mock 示意"""
        raise NotImplementedError

# ── 3. 调查流 Agent ──

class LogAnalystAgent(Agent):
    """模拟：从日志中提取异常信息"""
    def run(self, task: str, context: dict) -> dict:
        # 实际实现会调用日志系统 API + LLM 分析
        return {
            "agent": self.name,
            "error_codes": ["TIMEOUT", "PARTITION_MISSING"],
            "failed_at": "2024-06-12T03:15:00Z",
            "raw_log_snippet": "Task X failed: partition 2024-06-11 not found in table Y"
        }

class MetadataInspectorAgent(Agent):
    """模拟：检查 Metastore 中的表/分区状态"""
    def run(self, task: str, context: dict) -> dict:
        log_result = context.get("log_analysis", {})
        table_name = "Y"  # 实际应从 log_result 中提取
        return {
            "agent": self.name,
            "table": table_name,
            "last_partition": "2024-06-10",
            "expected_partition": "2024-06-11",
            "partition_gap_detected": True
        }

# ── 4. 增强流 Agent ──

class ScriptGeneratorAgent(Agent):
    """模拟：基于调查结论生成修复脚本"""
    def run(self, task: str, context: dict) -> dict:
        inv_result = context.get("investigation_result")
        if not inv_result or inv_result.confidence < 0.6:
            return {"agent": self.name, "skip_reason": "low confidence, defer to human"}

        # 实际实现会用 LLM 生成针对具体系统的脚本
        return {
            "agent": self.name,
            "fix_script": f"-- Add missing partition for {inv_result.root_cause}\n"
                          f"ALTER TABLE Y ADD PARTITION (ds='2024-06-11') "
                          f"LOCATION 's3://bucket/Y/ds=2024-06-11/';",
            "estimated_risk": "low"
        }

# ── 5. 编排层 ──

class Orchestrator:
    def __init__(self):
        self.log_agent = LogAnalystAgent("log-analyst")
        self.meta_agent = MetadataInspectorAgent("meta-inspector")
        self.script_agent = ScriptGeneratorAgent("script-generator")

    def handle_request(self, request: str) -> dict:
        context = {"original_request": request}

        # ── Phase 1: Investigation ──
        print(f"[Orchestrator] Starting investigation for: {request}")

        log_result = self.log_agent.run(request, context)
        context["log_analysis"] = log_result

        meta_result = self.meta_agent.run(request, context)
        context["meta_analysis"] = meta_result

        # 综合调查结论
        inv_result = InvestigationResult(
            root_cause=f"Partition gap: {meta_result['last_partition']} -> "
                       f"{meta_result['expected_partition']} missing",
            evidence=[
                log_result["raw_log_snippet"],
                f"Metastore shows last_partition={meta_result['last_partition']}"
            ],
            confidence=0.85 if meta_result["partition_gap_detected"] else 0.3,
            needs_human=not meta_result["partition_gap_detected"]
        )
        context["investigation_result"] = inv_result

        # ── 决策门：置信度不够则升级人工 ──
        if inv_result.needs_human or inv_result.confidence < 0.5:
            print("[Orchestrator] Low confidence — escalating to human.")
            return {
                "status": "escalated",
                "investigation": inv_result,
                "reason": "confidence below threshold"
            }

        # ── Phase 2: Enhancement ──
        print("[Orchestrator] Confidence OK — proceeding to enhancement.")
        script_result = self.script_agent.run(request, context)
        context["script_result"] = script_result

        enhancement = EnhancementPlan(
            actions=["add_missing_partition"],
            generated_script=script_result["fix_script"],
            doc_update=f"Root cause: {inv_result.root_cause}. Fix: add partition."
        )

        return {
            "status": "resolved_by_agent",
            "investigation": inv_result,
            "enhancement": enhancement,
            "requires_human_review": script_result.get("estimated_risk") != "low"
        }

# ── 6. 运行 ──

if __name__ == "__main__":
    orch = Orchestrator()
    result = orch.handle_request(
        "数据表 Y 的 6月12日 ETL 任务失败，请排查并修复"
    )
    print(json.dumps(result, default=lambda o: o.__dict__, indent=2, ensure_ascii=False))

{
  "status": "resolved_by_agent",
  "investigation": {
    "root_cause": "Partition gap: 2024-06-10 -> 2024-06-11 missing",
    "evidence": ["...log snippet...", "Metastore shows last_partition=2024-06-10"],
    "confidence": 0.85,
    "needs_human": false
  },
  "enhancement": {
    "actions": ["add_missing_partition"],
    "generated_script": "ALTER TABLE Y ADD PARTITION ...",
    "doc_update": "Root cause: Partition gap..."
  },
  "requires_human_review": false
}