LangChain 에이전트 무한루프·툴폭주 차단 7가지

import time
from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

prompt = ChatPromptTemplate.from_messages([
    ("system", "You are a careful agent. Stop when you have enough."),
    ("human", "{input}"),
    ("placeholder", "{agent_scratchpad}"),
])

agent = create_tool_calling_agent(llm, tools=[], prompt=prompt)

executor = AgentExecutor(
    agent=agent,
    tools=[],
    verbose=True,
    max_iterations=6,          # 반복 상한
    max_execution_time=12.0,   # 초 단위 시간 상한
    early_stopping_method="force",  # 강제 종료
)

result = executor.invoke({"input": "Solve X"})
print(result["output"])

import json
import hashlib
from collections import deque

class LoopGuard:
    def __init__(self, window=5, repeat_threshold=3):
        self.window = window
        self.repeat_threshold = repeat_threshold
        self.recent = deque(maxlen=window)

    def signature(self, tool_name: str, tool_args: dict) -> str:
        payload = json.dumps({"name": tool_name, "args": tool_args}, sort_keys=True)
        return hashlib.sha256(payload.encode()).hexdigest()

    def seen_too_much(self, tool_name: str, tool_args: dict) -> bool:
        sig = self.signature(tool_name, tool_args)
        self.recent.append(sig)
        return self.recent.count(sig) >= self.repeat_threshold

loop_guard = LoopGuard(window=6, repeat_threshold=2)

# tool 호출 직전
# if loop_guard.seen_too_much(name, args):
#     raise RuntimeError("Loop detected: repeated tool call")

import time
from dataclasses import dataclass

@dataclass
class ToolPolicy:
    max_calls: int
    cooldown_sec: float

class ToolGate:
    def __init__(self, policies: dict[str, ToolPolicy]):
        self.policies = policies
        self.calls = {k: 0 for k in policies.keys()}
        self.last_called_at = {k: 0.0 for k in policies.keys()}

    def check(self, tool_name: str):
        if tool_name not in self.policies:
            raise PermissionError(f"Tool not allowed: {tool_name}")

        p = self.policies[tool_name]
        now = time.time()

        if self.calls[tool_name] >= p.max_calls:
            raise RuntimeError(f"Tool quota exceeded: {tool_name}")

        if now - self.last_called_at[tool_name] < p.cooldown_sec:
            raise RuntimeError(f"Tool cooldown active: {tool_name}")

        self.calls[tool_name] += 1
        self.last_called_at[tool_name] = now

# 예: 검색은 최대 3회, 1초 쿨다운
# DB는 최대 2회, 쿨다운 없음
gate = ToolGate({
    "web_search": ToolPolicy(max_calls=3, cooldown_sec=1.0),
    "db_query": ToolPolicy(max_calls=2, cooldown_sec=0.0),
})

# tool wrapper 내부에서
# gate.check(tool_name)
# return tool.invoke(args)

from enum import Enum

class Stage(str, Enum):
    PLAN = "plan"
    EXECUTE = "execute"
    SYNTHESIZE = "synthesize"

state = {"stage": Stage.PLAN, "plan": None, "observations": []}

def plan_step(user_input: str):
    # LLM에게 "툴은 최대 2개만" 같은 제약을 강제
    return {"tools": ["web_search"], "queries": [user_input]}

def execute_step(plan: dict):
    obs = []
    for q in plan["queries"]:
        # web_search(q)
        obs.append({"q": q, "result": "..."})
    return obs

def synthesize_step(user_input: str, observations: list[dict]):
    # 관찰 기반으로 답변 생성 후 종료
    return "final answer"

user_input = "..."

state["plan"] = plan_step(user_input)
state["stage"] = Stage.EXECUTE

state["observations"] = execute_step(state["plan"])
state["stage"] = Stage.SYNTHESIZE

final = synthesize_step(user_input, state["observations"])
print(final)

from pydantic import BaseModel, Field, ValidationError

class SearchArgs(BaseModel):
    query: str = Field(min_length=2, max_length=200)
    top_k: int = Field(default=5, ge=1, le=10)

def safe_web_search(raw_args: dict):
    try:
        args = SearchArgs.model_validate(raw_args)
    except ValidationError as e:
        # 재시도 루프를 만들지 말고 즉시 관찰로 남기고 종료/대체
        raise ValueError(f"Invalid tool args: {e}")

    # 실제 검색 수행
    return {"items": [], "query": args.query, "top_k": args.top_k}

import asyncio

async def run_with_timeout(coro, timeout_sec: float):
    return await asyncio.wait_for(coro, timeout=timeout_sec)

# 예: 에이전트 실행을 10초로 제한
# result = await run_with_timeout(executor.ainvoke({"input": "..."}), 10.0)

from langchain_core.callbacks import BaseCallbackHandler

class GuardrailLogger(BaseCallbackHandler):
    def __init__(self):
        self.tool_calls = []

    def on_tool_start(self, serialized, input_str, **kwargs):
        name = serialized.get("name") if isinstance(serialized, dict) else str(serialized)
        self.tool_calls.append(name)
        print(f"tool_start name={name} total_calls={len(self.tool_calls)}")

    def on_agent_finish(self, finish, **kwargs):
        print(f"agent_finish tool_calls={self.tool_calls}")

handler = GuardrailLogger()

executor = AgentExecutor(
    agent=agent,
    tools=[],
    callbacks=[handler],
    max_iterations=6,
    max_execution_time=12.0,
)