LangChain 멀티에이전트 루프 방지 - 상태머신·가드레일

from dataclasses import dataclass, field
import time

@dataclass
class Budget:
    max_steps: int = 30
    max_tool_calls: int = 20
    deadline_epoch: float = field(default_factory=lambda: time.time() + 60)

@dataclass
class RuntimeCounters:
    steps: int = 0
    tool_calls: int = 0

class BudgetExceeded(Exception):
    pass

def enforce_budget(budget: Budget, c: RuntimeCounters):
    now = time.time()
    if c.steps >= budget.max_steps:
        raise BudgetExceeded("max_steps exceeded")
    if c.tool_calls >= budget.max_tool_calls:
        raise BudgetExceeded("max_tool_calls exceeded")
    if now >= budget.deadline_epoch:
        raise BudgetExceeded("deadline exceeded")

import json
import hashlib
from dataclasses import dataclass

@dataclass
class TransitionRecord:
    phase: str
    action: str
    action_input: dict
    observation_key: str

def stable_hash(d: dict) -> str:
    s = json.dumps(d, sort_keys=True, ensure_ascii=False)
    return hashlib.sha256(s.encode("utf-8")).hexdigest()

def transition_fingerprint(r: TransitionRecord) -> str:
    payload = {
        "phase": r.phase,
        "action": r.action,
        "action_input_hash": stable_hash(r.action_input),
        "observation_key": r.observation_key,
    }
    return stable_hash(payload)

class LoopDetected(Exception):
    pass

def detect_loop(history: list[TransitionRecord], window: int = 8, threshold: int = 2):
    # 최근 window 내 동일 fingerprint가 threshold+1번 이상 나오면 루프로 판단
    recent = history[-window:]
    fps = [transition_fingerprint(r) for r in recent]
    for fp in set(fps):
        if fps.count(fp) >= (threshold + 1):
            raise LoopDetected(f"repeated transition detected: {fp}")

from typing import TypedDict, Any, Literal

class ToolResult(TypedDict, total=False):
    ok: bool
    result: Any
    error_type: Literal[
        "NETWORK",
        "RATE_LIMIT",
        "AUTH",
        "INVALID_INPUT",
        "NOT_FOUND",
        "POLICY",
        "UNKNOWN",
    ]
    retryable: bool
    hint: str

def classify_exception(e: Exception) -> ToolResult:
    msg = str(e)
    # 예시는 단순화. 실제로는 HTTP status, provider error code로 분기
    if "429" in msg:
        return {"ok": False, "error_type": "RATE_LIMIT", "retryable": True, "hint": "backoff"}
    if "401" in msg or "403" in msg:
        return {"ok": False, "error_type": "AUTH", "retryable": False, "hint": "check credentials"}
    return {"ok": False, "error_type": "UNKNOWN", "retryable": False, "hint": "inspect logs"}

import json

class VerificationFailed(Exception):
    pass

def verify_json(text: str) -> dict:
    try:
        obj = json.loads(text)
    except json.JSONDecodeError as e:
        raise VerificationFailed(f"invalid json: {e}")

    # 필수 키 체크(간단 예시)
    required = ["answer", "sources"]
    missing = [k for k in required if k not in obj]
    if missing:
        raise VerificationFailed(f"missing keys: {missing}")

    return obj

from dataclasses import dataclass, field
from typing import Any

@dataclass
class AgentState:
    phase: str = "PLAN"
    goal: str = ""
    plan: list[str] = field(default_factory=list)
    last_output: Any = None
    last_error_type: str | None = None
    history: list[TransitionRecord] = field(default_factory=list)
    counters: RuntimeCounters = field(default_factory=RuntimeCounters)


def step(state: AgentState, budget: Budget) -> AgentState:
    enforce_budget(budget, state.counters)
    state.counters.steps += 1

    if state.phase == "PLAN":
        # planner_agent(...) 호출 결과라고 가정
        state.plan = ["task1", "task2"]
        state.history.append(TransitionRecord("PLAN", "make_plan", {"goal": state.goal}, "ok"))
        state.phase = "EXECUTE"
        return state

    if state.phase == "EXECUTE":
        enforce_budget(budget, state.counters)
        state.counters.tool_calls += 1

        # tool 호출/실행 결과라고 가정
        tool_res: ToolResult = {"ok": False, "error_type": "RATE_LIMIT", "retryable": True, "hint": "backoff"}

        obs_key = "ok" if tool_res.get("ok") else f"err:{tool_res.get('error_type')}"
        state.history.append(TransitionRecord("EXECUTE", "call_tool", {"task": state.plan[0]}, obs_key))
        detect_loop(state.history)

        if tool_res.get("ok"):
            state.last_output = tool_res.get("result")
            state.last_error_type = None
            state.phase = "VERIFY"
        else:
            state.last_error_type = tool_res.get("error_type")
            if not tool_res.get("retryable"):
                state.phase = "ABORT"
            else:
                state.phase = "REPAIR"
        return state

    if state.phase == "VERIFY":
        # verify(...) 수행
        state.history.append(TransitionRecord("VERIFY", "verify", {"has_output": state.last_output is not None}, "ok"))
        detect_loop(state.history)
        state.phase = "DONE"
        return state

    if state.phase == "REPAIR":
        # repair_agent(...) 호출 결과를 반영
        # 예: 백오프 후 재시도, 파라미터 수정, 대체 툴 선택
        state.history.append(TransitionRecord("REPAIR", "backoff", {"error": state.last_error_type}, "ok"))
        detect_loop(state.history)
        state.phase = "EXECUTE"
        return state

    return state


def run(goal: str):
    state = AgentState(goal=goal)
    budget = Budget(max_steps=15, max_tool_calls=5)

    while state.phase not in ["DONE", "ABORT"]:
        state = step(state, budget)

    return state