Middleware(Custom middleware)

HanJu Han·2025년 12월 7일

제 5장: 에이전트의 제어권, 커스텀 미들웨어

1. 미들웨어란 무엇인가? (remind)

미들웨어는 에이전트가 생각하고 행동하는 흐름(Flow) 중간에 개입하는 코드입니다.

상상해 보세요. 여러분이 백화점의 매니저입니다. 직원이 손님을 응대할 때마다 옆에서 지켜보다가:
1. 대화 시작 전: "오늘 VIP 손님이니 정중하게 해"라고 귀띔하거나 (설정 주입)
2. 물건 건네기 전: "잠깐, 이 물건 재고가 맞는지 내가 먼저 확인할게"라고 검사하거나 (유효성 검사)
3. 실수했을 때: "다시 한번 정중하게 말해봐"라고 기회를 주는 것 (재시도)

이 모든 것이 미들웨어의 역할입니다. LangChain/LangGraph에서는 크게 두 가지 스타일의 훅(Hook, 갈고리)을 제공합니다.


2. 두 가지 스타일의 훅 (Hooks)

스타일 1: 노드 스타일 훅 (Node-style hooks)

  • 특징: 실행 흐름의 특정 지점(체크포인트)에서 순차적으로 실행됩니다.
  • 용도: 기록(Logging), 검사(Validation), 상태 변경.
  • 종류:
    • before_agent: 에이전트 시작 전 (딱 한 번)
    • before_model: 모델(LLM)을 호출하기 직전
    • after_model: 모델이 응답한 직후
    • after_agent: 에이전트 종료 후 (딱 한 번)

스타일 2: 랩 스타일 훅 (Wrap-style hooks)

  • 특징: 실행 자체를 감싸버립니다(Wrap). 핸들러를 호출할지, 말지, 몇 번 할지를 결정할 수 있습니다.
  • 용도: 재시도(Retry), 캐싱(Caching), 결과 조작.
  • 종류:
    • wrap_model_call: 모델 호출을 감쌈
    • wrap_tool_call: 도구(Tool) 호출을 감쌈

3. 시각화로 이해하는 실행 흐름 (Mermaid)

이커머스 에이전트가 사용자의 주문을 처리할 때 미들웨어가 어디서 작동하는지 살펴봅시다.

  • 분홍색 마름모: 노드 스타일 훅 (체크포인트)
  • 파란색 박스: 랩 스타일 훅 (감싸서 제어)

4. 실전 예제: 이커머스 'SmartShopper' 만들기

이제 코드를 작성해 봅시다. 우리는 두 가지 미들웨어를 만들 것입니다.

  1. 예산 지킴이 (Node-style): 사용자가 너무 비싼 물건을 사려 하거나 대화가 너무 길어지면 차단합니다.
  2. 끈기 있는 점원 (Wrap-style): 재고 조회 시스템이 불안정할 때, 자동으로 재시도합니다.

(1) 데코레이터(Decorator) 방식

가장 빠르고 간편하게 미들웨어를 만드는 방법입니다.

from langchain.agents.middleware import before_model, wrap_tool_call, AgentState, ModelRequest, ModelResponse
from langchain.messages import AIMessage
from langgraph.runtime import Runtime
from typing import Any, Callable
import time

# ---------------------------------------------------------
# 1. 노드 스타일 훅: 예산 및 대화 길이 제한 (Before Model)
# ---------------------------------------------------------
@before_model(can_jump_to=["end"]) # 흐름을 강제로 끝낼 수 있는 권한 부여
def budget_guardian(state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
    """
    모델이 호출되기 전에 실행됩니다.
    대화가 너무 길어지면(비용 문제) 상담을 강제 종료합니다.
    """
    # 이커머스 시나리오: 대화 턴 수가 10회를 넘어가면 상담 종료
    if len(state["messages"]) >= 10:
        print("🚫 [경고] 상담이 너무 길어졌습니다. 예산 보호를 위해 종료합니다.")
        return {
            "messages": [AIMessage(content="죄송합니다. 상담 시간이 초과되어 상담을 종료합니다. 고객센터로 문의해주세요.")],
            "jump_to": "end" # 에이전트 실행을 즉시 종료
        }
    
    print(f"✅ [확인] 현재 대화 턴 수: {len(state['messages'])} (정상)")
    return None

# ---------------------------------------------------------
# 2. 랩 스타일 훅: 재고 조회 재시도 로직 (Wrap Tool Call)
# ---------------------------------------------------------
@wrap_tool_call
def robust_inventory_check(
    request: ModelRequest, # 툴 호출 요청 정보
    handler: Callable[[ModelRequest], ModelResponse], # 실제 툴을 실행하는 함수
) -> ModelResponse:
    """
    툴 호출을 감싸서, 에러 발생 시 3번까지 재시도합니다.
    """
    max_retries = 3
    
    for attempt in range(max_retries):
        try:
            # 실제 툴(핸들러) 실행
            return handler(request)
        except Exception as e:
            print(f"⚠️ [재고 시스템 에러] 시도 {attempt + 1}/{max_retries} 실패: {e}")
            if attempt == max_retries - 1:
                # 마지막 시도도 실패하면 에러를 그대로 던짐
                raise e
            time.sleep(1) # 잠시 대기 후 재시도
            
    return handler(request) # 도달하지 않음

# ---------------------------------------------------------
# 에이전트 생성 (가상 코드)
# ---------------------------------------------------------
# agent = create_agent(
#     model="gpt-4o",
#     middleware=[budget_guardian, robust_inventory_check], # 미들웨어 등록
#     tools=[inventory_tool],
# )

(2) 클래스(Class) 방식

더 복잡하고 체계적인 관리가 필요할 때 사용합니다. 여러 훅을 하나의 클래스로 묶을 수 있습니다.

시나리오: '쇼핑 감사(Audit) 미들웨어'를 만들어 봅시다. 이 미들웨어는 모델 호출 전후를 모두 기록하고, 설정에 따라 동작을 바꿉니다.

from langchain.agents.middleware import AgentMiddleware

class ShoppingAuditMiddleware(AgentMiddleware):
    def __init__(self, verbose: bool = True):
        self.verbose = verbose

    # 1. 모델 호출 전: 고객의 의도 파악 로그
    def before_model(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
        if self.verbose:
            last_msg = state["messages"][-1]
            print(f"🛒 [Audit-Start] 고객의 마지막 메시지 처리 중: {last_msg.content[:20]}...")
        return None

    # 2. 모델 호출 후: 추천 상품 기록
    def after_model(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
        if self.verbose:
            ai_response = state["messages"][-1].content
            print(f"📦 [Audit-End] AI 응답 완료. 길이: {len(ai_response)}자")
        return None

# 사용법
# agent = create_agent(
#     model="gpt-4o",
#     middleware=[ShoppingAuditMiddleware(verbose=True)], # 클래스 인스턴스 주입
#     tools=[...],
# )

5. 언제 무엇을 써야 할까요?

구분데코레이터 (@)클래스 (Class)
복잡도단순함 (함수 하나)복잡함 (여러 메서드)
설정코드 내 하드코딩 필요__init__으로 설정값 주입 가능
훅 개수훅 1개당 함수 1개하나의 클래스에 before, after 등 여러 훅 통합 가능
추천 상황빠른 프로토타이핑, 단순 로깅라이브러리 제작, 복잡한 상태 관리, 재사용성 중시

6. 핵심 요약

  1. 미들웨어는 안전벨트입니다. 특히 이커머스처럼 결제나 개인정보가 오가는 시스템에서는 before_model을 통해 민감한 정보가 유출되거나 잘못된 주문이 들어가는 것을 사전에 차단해야 합니다.
  2. Wrap 스타일은 복원력을 높입니다. 외부 API(재고, 배송 조회)는 언제든 실패할 수 있습니다. wrap_tool_call로 재시도 로직을 넣으면 에이전트가 훨씬 똑똑해 보입니다.
  3. jump_to="end"를 기억하세요. 노드 스타일 훅에서 문제가 발견되면, 굳이 모델을 호출하여 돈을 낭비하지 말고 즉시 종료시키는 것이 비용 절감의 핵심입니다.

지난 장에서 미들웨어의 기본 개념을 익혔다면, 이번 장에서는 "미들웨어의 심화 기술"을 다룹니다. 단순히 로깅을 남기는 수준을 넘어, 에이전트의 기억(State)을 확장하고, 실행 순서(Order)를 제어하며, 필요하다면 실행 흐름을 강제로 변경(Jump)하는 고급 기법들입니다.

이 기술들은 복잡한 비즈니스 로직을 가진 이커머스 에이전트를 구축할 때 필수적입니다. 자, 깊이 있게 들어가 봅시다.


제 6장: 미들웨어 심화 - 상태 확장과 제어 흐름

1. 커스텀 상태 스키마 (Custom State Schema)

기본적으로 에이전트는 messages(대화 내역)라는 상태를 가지고 있습니다. 하지만 쇼핑몰 상담원이라면 이것만으로는 부족합니다.
"이 고객이 VIP인가?", "쿠폰 조회는 몇 번이나 시도했나?", "현재 장바구니 금액은 얼마인가?" 같은 추가적인 정보가 필요합니다.

미들웨어는 에이전트의 기본 상태(State)에 사용자 정의 필드를 추가하여, 훅(Hook)끼리 데이터를 공유하거나 실행 흐름을 제어하는 데 사용할 수 있습니다.

이커머스 시나리오: '쿠폰 남용 방지 시스템'

고객이 상담 도중 쿠폰 조회를 3회 이상 시도하면, 더 이상 조회를 못 하게 막는 기능을 만들어 봅시다. 이를 위해 coupon_check_count라는 상태가 필요합니다.

코드 구현

from langchain.agents import create_agent
from langchain.messages import HumanMessage, AIMessage
from langchain.agents.middleware import AgentState, before_model, after_model
from typing_extensions import NotRequired
from typing import Any
from langgraph.runtime import Runtime

# 1. 커스텀 상태 정의 (기억 확장)
class ShoppingState(AgentState):
    # 기본 messages 외에 추가할 필드들
    coupon_check_count: NotRequired[int]  # 쿠폰 조회 횟수
    user_tier: NotRequired[str]           # 고객 등급 (예: 'VIP', 'Basic')

# 2. Before 훅: 횟수 제한 검사
@before_model(state_schema=ShoppingState, can_jump_to=["end"])
def check_coupon_limit(state: ShoppingState, runtime: Runtime) -> dict[str, Any] | None:
    # 상태에서 현재 횟수 가져오기 (기본값 0)
    count = state.get("coupon_check_count", 0)
    tier = state.get("user_tier", "Basic")
    
    print(f"🔍 [검사] 고객 등급: {tier}, 현재 조회 횟수: {count}")

    # 일반 고객이 3회 이상 조회 시 차단
    if tier == "Basic" and count >= 3:
        return {
            "messages": [AIMessage("죄송합니다. 일반 회원의 쿠폰 조회 한도를 초과했습니다.")],
            "jump_to": "end" # 상담 강제 종료
        }
    return None

# 3. After 훅: 횟수 증가 (카운터)
@after_model(state_schema=ShoppingState)
def increment_coupon_counter(state: ShoppingState, runtime: Runtime) -> dict[str, Any] | None:
    # 모델이 응답한 후 카운트 1 증가
    current_count = state.get("coupon_check_count", 0)
    return {"coupon_check_count": current_count + 1}

# 4. 에이전트 생성
agent = create_agent(
    model="gpt-4o",
    middleware=[check_coupon_limit, increment_coupon_counter],
    tools=[],
)

# 5. 실행 (초기 상태 주입)
result = agent.invoke({
    "messages": [HumanMessage("쿠폰 좀 찾아줘")],
    "coupon_check_count": 2, # 이미 2번 조회했다고 가정
    "user_tier": "Basic",
})

2. 미들웨어 실행 순서 (Execution Order)

여러 개의 미들웨어를 동시에 사용할 때, 어떤 순서로 실행되는지 이해하는 것은 매우 중요합니다. 이를 '양파 껍질(Onion)' 구조라고 생각하면 쉽습니다.

  • Before 훅: 리스트의 앞에서 뒤로 (First to Last) 실행됩니다. (진입)
  • After 훅: 리스트의 뒤에서 앞으로 (Last to First) 실행됩니다. (탈출)
  • Wrap 훅: 첫 번째 미들웨어가 나머지를 감싸는 형태입니다.

시각화 (Mermaid)

미들웨어 리스트가 [보안검사, 로깅, 데이터변환] 순서로 등록되어 있다고 가정해 봅시다.

핵심 규칙:
1. 가장 먼저 등록된 미들웨어(M1)가 가장 먼저 요청을 받고, 가장 나중에 응답을 내보냅니다.
2. before 로직은 순서대로, after 로직은 역순으로 실행됩니다.


3. 에이전트 점프 (Agent Jumps)

미들웨어는 단순히 지켜보는 것뿐만 아니라, 흐름을 납치(Hijack)할 수 있습니다. 특정 조건에서 에이전트의 실행 경로를 바꿀 때 jump_to를 사용합니다.

사용 가능한 점프 타겟

  • end: 에이전트 실행을 즉시 종료합니다. (가장 많이 사용됨)
  • tools: 모델 호출을 건너뛰고 바로 도구 실행 단계로 갑니다.
  • model: 도구 실행 등을 건너뛰고 다시 모델 호출 단계로 갑니다.

이커머스 시나리오: '블랙리스트 키워드 차단'

고객이 경쟁사 이름이나 부적절한 단어를 언급하면, 모델에게 물어볼 필요도 없이 즉시 대화를 차단해야 합니다.

from langchain.agents.middleware import after_model, hook_config

@after_model
@hook_config(can_jump_to=["end"]) # 점프 권한 설정 필수
def block_competitor_mention(state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
    last_message = state["messages"][-1]
    
    # 금지어 목록
    banned_words = ["경쟁사A", "나쁜말"]
    
    if any(word in last_message.content for word in banned_words):
        print("🚫 [차단] 금지된 키워드가 감지되었습니다.")
        return {
            # 사용자에게 보여줄 메시지로 덮어쓰기
            "messages": [AIMessage("죄송합니다. 해당 주제에 대해서는 답변해 드릴 수 없습니다.")],
            # 즉시 종료
            "jump_to": "end"
        }
    return None

4. 핵심 요약

  1. 단일 책임 원칙 (Keep it Focused):
    • 하나의 미들웨어는 하나의 일만 하세요. '로깅'과 '보안 검사'를 하나의 함수에 섞지 마세요.
  2. 우아한 에러 처리 (Handle Errors Gracefully):
    • 미들웨어에서 에러가 나면 에이전트 전체가 멈춥니다. try-except 구문으로 미들웨어 내부 에러가 메인 로직을 방해하지 않도록 하세요.
  3. 적절한 훅 타입 선택:
    • 단순 기록/검사 -> Node-style (before, after)
    • 재시도/캐싱/흐름 제어 -> Wrap-style
  4. 순서가 생명입니다:
    • 보안이나 인증 같은 중요한 미들웨어는 리스트의 맨 앞에 두세요. 그래야 가장 먼저 검사하고, 가장 나중에 결과를 최종 승인할 수 있습니다.
  5. 문서화:
    • CustomState에 어떤 필드가 추가되는지 팀원들이 알 수 있도록 명확히 주석을 남기세요.
profile
시리즈를 기반으로 작성하였습니다.

0개의 댓글