LangChain Agent 무한루프·토큰폭탄 차단 5팁

import hashlib
import json
from collections import deque

class LoopDetector:
    def __init__(self, window=8, threshold=3):
        self.window = window
        self.threshold = threshold
        self.recent = deque(maxlen=window)

    def _sig(self, tool_name: str, tool_input: dict) -> str:
        payload = json.dumps({"tool": tool_name, "input": tool_input}, sort_keys=True)
        return hashlib.sha256(payload.encode("utf-8")).hexdigest()

    def check_and_record(self, tool_name: str, tool_input: dict) -> None:
        sig = self._sig(tool_name, tool_input)
        self.recent.append(sig)
        if self.recent.count(sig) >= self.threshold:
            raise RuntimeError("Loop detected: repeated tool call pattern")

import time

class Budget:
    def __init__(self, max_steps=12, max_tool_calls=8, max_seconds=25):
        self.max_steps = max_steps
        self.max_tool_calls = max_tool_calls
        self.max_seconds = max_seconds
        self.steps = 0
        self.tool_calls = 0
        self.started = time.time()

    def step(self):
        self.steps += 1
        if self.steps > self.max_steps:
            raise TimeoutError("Budget exceeded: max steps")
        if time.time() - self.started > self.max_seconds:
            raise TimeoutError("Budget exceeded: max seconds")

    def tool_call(self):
        self.tool_calls += 1
        if self.tool_calls > self.max_tool_calls:
            raise TimeoutError("Budget exceeded: max tool calls")

def truncate(text: str, limit: int = 2000) -> str:
    if len(text) <= limit:
        return text
    return text[:limit] + "\n[TRUNCATED]"


def normalize_tool_output(output: str) -> dict:
    # 운영에서는 JSON schema로 정규화하는 편이 더 안전합니다.
    return {
        "preview": truncate(output, 1200),
        "length": len(output),
    }

from dataclasses import dataclass, field

@dataclass
class AgentState:
    facts: list[str] = field(default_factory=list)
    todos: list[str] = field(default_factory=list)
    last_tool_result_preview: str | None = None

    def to_prompt_block(self) -> str:
        facts = "\n".join(f"- {x}" for x in self.facts[-10:])
        todos = "\n".join(f"- {x}" for x in self.todos[-10:])
        return (
            "[STATE]\n"
            f"Facts:\n{facts}\n\n"
            f"Todos:\n{todos}\n\n"
            f"LastResultPreview:\n{self.last_tool_result_preview or ''}\n"
        )

import asyncio

class ToolTimeout(Exception):
    pass

async def run_tool_with_timeout(tool_coro, timeout_s: float = 8.0):
    try:
        return await asyncio.wait_for(tool_coro, timeout=timeout_s)
    except asyncio.TimeoutError as e:
        raise ToolTimeout("Tool call timed out") from e

class AgentAbort(Exception):
    def __init__(self, user_message: str, debug_reason: str):
        super().__init__(debug_reason)
        self.user_message = user_message
        self.debug_reason = debug_reason


def abort_with_partial(partial: str, reason: str) -> dict:
    return {
        "status": "aborted",
        "message": "작업이 길어져 중단했습니다. 아래는 현재까지의 결과입니다.",
        "partial_result": partial,
        "reason": reason,
    }

async def agent_loop(agent, tools, state: AgentState):
    budget = Budget(max_steps=12, max_tool_calls=8, max_seconds=25)
    loop_detector = LoopDetector(window=8, threshold=3)

    partial = ""

    while True:
        budget.step()

        prompt = state.to_prompt_block() + "\n[USER]\n" + agent.user_input
        action = await agent.plan(prompt)  # tool_name, tool_input, or final

        if action.get("final"):
            return {"status": "ok", "result": action["final"]}

        tool_name = action["tool"]
        tool_input = action.get("input", {})

        loop_detector.check_and_record(tool_name, tool_input)
        budget.tool_call()

        tool = tools[tool_name]
        try:
            raw = await run_tool_with_timeout(tool(**tool_input), timeout_s=8.0)
        except Exception as e:
            # 실패를 모델에게 장문으로 주지 말고, 짧은 상태로 반영
            state.last_tool_result_preview = f"ToolError: {type(e).__name__}"
            state.todos.append("다른 접근(쿼리 축소/캐시/추가 질문) 시도")
            continue

        normalized = normalize_tool_output(str(raw))
        state.last_tool_result_preview = normalized["preview"]
        partial = state.last_tool_result_preview

        # 필요하면 facts/todos를 갱신하는 별도 로직을 둡니다.
        state.facts.append(f"Used tool {tool_name}, got {normalized['length']} chars")