LangGraph로 AutoGPT 다중에이전트 상태머신 구축

from __future__ import annotations

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

# LangGraph
from langgraph.graph import StateGraph, END


@dataclass
class Task:
    id: str
    description: str
    status: str = "todo"  # todo | doing | done | failed
    attempts: int = 0


@dataclass
class AgentState:
    goal: str
    tasks: List[Task] = field(default_factory=list)
    current_task_id: Optional[str] = None
    artifacts: Dict[str, Any] = field(default_factory=dict)
    observations: List[Dict[str, Any]] = field(default_factory=list)
    messages: List[Dict[str, str]] = field(default_factory=list)
    loop_guard: Dict[str, Any] = field(default_factory=lambda: {
        "steps": 0,
        "max_steps": 30,
        "recent_nodes": [],
        "max_same_node": 5,
    })
    errors: List[Dict[str, Any]] = field(default_factory=list)

import json
import time


def call_llm(messages: List[Dict[str, str]]) -> str:
    # 예시용 더미. 실제로는 모델 API 호출.
    return "{}"


def planner_node(state: AgentState) -> AgentState:
    state.loop_guard["steps"] += 1

    system = {
        "role": "system",
        "content": "You are a planner. Decompose the goal into 3-7 actionable tasks.",
    }
    user = {"role": "user", "content": state.goal}

    raw = call_llm([system, user])

    # 예: 모델이 JSON 배열을 준다고 가정
    try:
        items = json.loads(raw) if raw.strip() else []
    except Exception as e:
        state.errors.append({"where": "planner", "error": str(e), "raw": raw})
        items = []

    if not items:
        # 최소 안전장치
        items = [
            {"id": "t1", "description": "Clarify requirements and constraints"},
            {"id": "t2", "description": "Draft solution outline"},
            {"id": "t3", "description": "Produce final output"},
        ]

    state.tasks = [Task(id=i["id"], description=i["description"]) for i in items]
    state.messages = [{"role": "system", "content": "Working on tasks."}]
    return state


def router_node(state: AgentState) -> AgentState:
    state.loop_guard["steps"] += 1

    # 다음 todo 선택
    for t in state.tasks:
        if t.status == "todo":
            t.status = "doing"
            state.current_task_id = t.id
            return state

    # 모두 끝났으면 finalizer로
    state.current_task_id = None
    return state


def executor_node(state: AgentState) -> AgentState:
    state.loop_guard["steps"] += 1

    task = next((t for t in state.tasks if t.id == state.current_task_id), None)
    if not task:
        state.errors.append({"where": "executor", "error": "no current task"})
        return state

    task.attempts += 1

    # 실제론 여기서 툴 호출, 검색, 코드 실행 등을 수행
    # 예시로는 관찰 로그만 추가
    obs = {
        "task_id": task.id,
        "ts": time.time(),
        "result": f"Executed: {task.description}",
    }
    state.observations.append(obs)

    # 산출물은 요약 형태로만 누적
    state.artifacts[task.id] = {"summary": obs["result"]}

    return state


def critic_node(state: AgentState) -> AgentState:
    state.loop_guard["steps"] += 1

    task = next((t for t in state.tasks if t.id == state.current_task_id), None)
    if not task:
        return state

    # 간단 검증 규칙: attempts 2회 초과면 실패 처리
    if task.attempts >= 2 and task.status != "done":
        task.status = "failed"
        state.errors.append({
            "where": "critic",
            "task_id": task.id,
            "error": "max attempts exceeded",
        })
        return state

    # 예시에서는 항상 성공 처리
    task.status = "done"
    return state


def finalizer_node(state: AgentState) -> AgentState:
    state.loop_guard["steps"] += 1

    # 최종 결과를 artifacts에서 조립
    done = [t for t in state.tasks if t.status == "done"]
    failed = [t for t in state.tasks if t.status == "failed"]

    state.artifacts["final"] = {
        "done": [t.id for t in done],
        "failed": [t.id for t in failed],
        "notes": "Finalized output.",
    }
    return state

def guard_or_next(state: AgentState, default_next: str) -> str:
    lg = state.loop_guard

    if lg["steps"] >= lg["max_steps"]:
        return "finalizer"

    # 최근 노드 반복 감지(운영에서 꽤 유용)
    recent = lg["recent_nodes"]
    recent.append(default_next)
    lg["recent_nodes"] = recent[-20:]

    if len(recent) >= lg["max_same_node"]:
        tail = recent[-lg["max_same_node"]:]
        if all(x == tail[0] for x in tail):
            state.errors.append({
                "where": "loop_guard",
                "error": "same node repeated",
                "node": tail[0],
            })
            return "finalizer"

    return default_next


def route_after_router(state: AgentState) -> str:
    if state.current_task_id is None:
        return guard_or_next(state, "finalizer")
    return guard_or_next(state, "executor")


def route_after_critic(state: AgentState) -> str:
    # 현재 task가 failed면 다음 task로 이동
    task = next((t for t in state.tasks if t.id == state.current_task_id), None)
    if task and task.status in ["done", "failed"]:
        return guard_or_next(state, "router")

    # 그 외는 재시도
    return guard_or_next(state, "executor")

def build_graph():
    g = StateGraph(AgentState)

    g.add_node("planner", planner_node)
    g.add_node("router", router_node)
    g.add_node("executor", executor_node)
    g.add_node("critic", critic_node)
    g.add_node("finalizer", finalizer_node)

    g.set_entry_point("planner")

    g.add_edge("planner", "router")

    g.add_conditional_edges(
        "router",
        route_after_router,
        {
            "executor": "executor",
            "finalizer": "finalizer",
        },
    )

    g.add_edge("executor", "critic")

    g.add_conditional_edges(
        "critic",
        route_after_critic,
        {
            "router": "router",
            "executor": "executor",
            "finalizer": "finalizer",
        },
    )

    g.add_edge("finalizer", END)

    return g.compile()


if __name__ == "__main__":
    app = build_graph()
    init = AgentState(goal="Design a multi-agent AutoGPT system with safety guards")
    out = app.invoke(init)
    print(out.artifacts.get("final"))