커스텀 미들웨어 만들기

커스텀 미들웨어 만들기 (Custom middleware)

에이전트 실행 흐름의 특정 지점마다 훅을 구현해서 동작을 가로채는, 나만의 커스텀 미들웨어를 만드는 방법을 살펴볼게요. 에이전트가 모델을 부르기 전후, 도구를 실행하기 전후 등 원하는 시점에 코드를 끼워 넣을 수 있죠. 미리 만들어 둔 미들웨어가 컵라면이라면, 커스텀 미들웨어는 직접 국물을 내는 요리라고 할 수 있어요.

출처: LangChain 공식 문서 — middleware-custom

훅(Hooks)이 뭔가요?

미들웨어는 에이전트 실행을 가로채기 위해 두 가지 스타일의 훅을 제공해요.

• 노드 스타일 훅 (Node-style hooks): 특정 실행 지점에서 순차적으로 실행됩니다.
• 래핑 스타일 훅 (Wrap-style hooks): 각 모델 호출이나 도구 호출을 감싸서 실행되죠.

노드 스타일 훅

특정 실행 지점에서 순차적으로 실행돼요. 로깅, 검증, 상태 업데이트에 주로 쓰입니다. 필요에 따라 어떤 훅을 쓸지 고르면 돼요.

훅	실행 시점
before_agent	에이전트 시작 전 (호출당 1회)
before_model	각 모델 호출 전
after_model	각 모델 응답 후
after_agent	에이전트 완료 후 (호출당 1회)

랩 스타일 훅은 각 호출을 감싸서 실행에 대한 제어권을 줍니다.

훅	실행 시점
wrap_model_call	각 모델 호출을 감싸서
wrap_tool_call	각 도구 호출을 감싸서

데코레이터로 만드는 예시

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

@before_model(can_jump_to=["end"])
def check_message_limit(state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
    if len(state["messages"]) >= 50:
        return {
            "messages": [AIMessage("Conversation limit reached.")],
            "jump_to": "end"
        }
    return None

@after_model
def log_response(state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
    print(f"Model returned: {state['messages'][-1].content}")
    return None

클래스로 만드는 예시

from langchain.agents.middleware import AgentMiddleware, AgentState, hook_config
from langchain.messages import AIMessage
from langgraph.runtime import Runtime
from typing import Any

class MessageLimitMiddleware(AgentMiddleware):
    def __init__(self, max_messages: int = 50):
        super().__init__()
        self.max_messages = max_messages

    @hook_config(can_jump_to=["end"])
    def before_model(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
        if len(state["messages"]) >= self.max_messages:
            return {
                "messages": [AIMessage("Conversation limit reached.")],
                "jump_to": "end"
            }
        return None

    def after_model(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
        print(f"Model returned: {state['messages'][-1].content}")
        return None

랩 스타일 훅 (Wrap-style hooks)

실행을 가로채고 핸들러가 언제 호출되는지 제어해요. 재시도(retry), 캐싱, 변환(transformation)에 주로 쓰입니다. 핸들러를 0번 호출할지(short-circuit), 1번 호출할지(정상 흐름), 여러 번 호출할지(재시도 로직)를 직접 결정할 수 있어요.

• wrap_model_call — 각 모델 호출을 감쌈
• wrap_tool_call — 각 도구 호출을 감쌈

from langchain.agents.middleware import wrap_model_call, ModelRequest, ModelResponse
from typing import Callable

@wrap_model_call
def retry_model(
    request: ModelRequest,
    handler: Callable[[ModelRequest], ModelResponse],
) -> ModelResponse:
    for attempt in range(3):
        try:
            return handler(request)
        except Exception as e:
            if attempt == 2:
                raise
            print(f"Retry {attempt + 1}/3 after error: {e}")

상태 업데이트 (State updates)

노드 스타일과 랩 스타일 훅 모두 에이전트 상태를 업데이트할 수 있는데, 방식이 달라요.

• 노드 스타일 훅 (before_agent, before_model, after_model, after_agent): 딕셔너리를 직접 반환해요. 이 딕셔너리는 그래프의 리듀서(reducers)를 통해 에이전트 상태에 반영됩니다.
• 랩 스타일 훅 (wrap_model_call, wrap_tool_call): 모델 호출의 경우 ExtendedModelResponse에 Command를 담아 상태 업데이트를 주입하고, 도구 호출의 경우 Command를 직접 반환해요. 요약 트리거 지점, 사용량 메타데이터, 요청이나 응답에서 계산한 커스텀 필드처럼 모델·도구 호출 중 실행되는 로직을 기반으로 상태를 추적하거나 업데이트할 때 사용합니다.

from langchain.agents.middleware import after_model, AgentState
from langgraph.runtime import Runtime
from typing import Any
from typing_extensions import NotRequired


class TrackingState(AgentState):
    model_call_count: NotRequired[int]


@after_model(state_schema=TrackingState)
def increment_after_model(state: TrackingState, runtime: Runtime) -> dict[str, Any] | None:
    return {"model_call_count": state.get("model_call_count", 0) + 1}

래핑 스타일 훅 쪽은 이렇게 Command를 사용해요.

from typing import Callable
from langchain.agents.middleware import (
    wrap_model_call,
    ModelRequest,
    ModelResponse,
    AgentState,
    ExtendedModelResponse
)
from langgraph.types import Command
from typing_extensions import NotRequired

class UsageTrackingState(AgentState):
    """Agent state with token usage tracking."""

    last_model_call_tokens: NotRequired[int]


@wrap_model_call(state_schema=UsageTrackingState)
def track_usage(
    request: ModelRequest,
    handler: Callable[[ModelRequest], ModelResponse],
) -> ExtendedModelResponse:
    response = handler(request)
    return ExtendedModelResponse(
        model_response=response,
        command=Command(update={"last_model_call_tokens": 150}),
    )

Command는 그래프의 리듀서를 거치므로 업데이트가 올바르게 적용되고, 메시지는 기존 상태를 덮어쓰지 않고 덧붙여져요(additive).

여러 미들웨어 합성 (Composition with multiple middleware)

여러 미들웨어 레이어가 ExtendedModelResponse를 반환하면 그 Command들이 합성됩니다.

• 리듀서를 통한 적용: 각 Command는 별도의 상태 업데이트가 됩니다. 메시지의 경우 덧붙이는 방식이죠.
• 충돌 시 바깥쪽이 승리: 리듀서가 없는 상태 필드의 경우 안쪽부터 적용한 뒤 바깥쪽이 마지막에 적용되어, 충돌하는 키에서는 가장 바깥쪽 미들웨어의 값이 우선합니다.
• Retry-safe: 바깥쪽 미들웨어가 handler()를 여러 번 호출하는 로직(예: 재시도)을 가진 경우, 이전 호출의 Command는 폐기됩니다.

from typing import Annotated, Callable

from langchain.agents.middleware import (
    AgentMiddleware,
    AgentState,
    ExtendedModelResponse,
    ModelRequest,
    ModelResponse,
)
from langchain.messages import SystemMessage
from langgraph.types import Command
from typing_extensions import NotRequired


def _last_wins(_a: str, b: str) -> str:
    """Reducer: last writer wins (outer overwrites inner)."""
    return b

class CustomMiddlewareState(AgentState):
    """Agent state: trace_layer uses last-wins (outer wins), messages use additive reducer."""

    # Non-reducer field with last-wins: both middleware write; outermost value wins
    trace_layer: NotRequired[Annotated[str, _last_wins]]


class OuterMiddleware(AgentMiddleware):
    def wrap_model_call(
        self,
        request: ModelRequest,
        handler: Callable[[ModelRequest], ModelResponse],
    ) -> ExtendedModelResponse:
        response = handler(request)
        return ExtendedModelResponse(
            model_response=response,
            command=Command(update={
                "trace_layer": "outer",
                "messages": [SystemMessage(content="[Outer ran]")],
            }),
        )


class InnerMiddleware(AgentMiddleware):
    """Adds trace_layer and message. Outer adds to same keys; trace_layer: outer wins, messages: additive."""

    def wrap_model_call(
        self,
        request: ModelRequest,
        handler: Callable[[ModelRequest], ModelResponse],
    ):
        response = handler(request)
        return ExtendedModelResponse(
            model_response=response,
            command=Command(update={
                "trace_layer": "inner",
                "messages": [SystemMessage(content="[Inner ran]")],
            }),
        )

미들웨어 만들기 (Create middleware)

미들웨어를 만드는 방법은 두 가지예요.

데코레이터 기반 미들웨어

단일 훅 미들웨어에 빠르고 간단한 방식입니다. 개별 함수를 데코레이터로 감싸죠. 사용 가능한 데코레이터:

• 노드 스타일: @before_agent, @before_model, @after_model, @after_agent
• 랩 스타일: @wrap_model_call, @wrap_tool_call
• 편의용: @dynamic_prompt

from langchain.agents.middleware import (
    before_model,
    wrap_model_call,
    AgentState,
    ModelRequest,
    ModelResponse,
)
from langchain.agents import create_agent
from langgraph.runtime import Runtime
from typing import Any, Callable


@before_model
def log_before_model(state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
    print(f"About to call model with {len(state['messages'])} messages")
    return None

@wrap_model_call
def retry_model(
    request: ModelRequest,
    handler: Callable[[ModelRequest], ModelResponse],
) -> ModelResponse:
    for attempt in range(3):
        try:
            return handler(request)
        except Exception as e:
            if attempt == 2:
                raise
            print(f"Retry {attempt + 1}/3 after error: {e}")

agent = create_agent(
    model="gpt-5.5",
    middleware=[log_before_model, retry_model],
    tools=[...],
)

데코레이터를 쓸 때: 단일 훅만 필요할 때, 복잡한 설정이 필요 없을 때, 빠른 프로토타이핑을 할 때 좋아요.

클래스 기반 미들웨어

여러 훅이나 설정이 필요한 복잡한 미들웨어에 더 강력한 방식입니다. 같은 훅에 대해 동기·비동기 구현을 모두 정의하거나, 하나의 미들웨어에서 여러 훅을 조합할 때 클래스를 써요.

AgentMiddleware 서브클래스는 에이전트 팩토리가 컴파일 시점에 읽어들이는 세 가지 클래스 속성을 선언할 수 있습니다:

• state_schema — 커스텀 필드로 에이전트 상태를 확장합니다 (Custom state schema 참고).
• tools — 미들웨어와 함께 제공되는 추가 도구를 등록합니다 (예: to-do 리스트 미들웨어의 write_todos).
• transformers — 스코프 인지형 스트림 변압기 팩토리를 등록합니다 (Custom stream transformers 참고).

from langchain.agents.middleware import (
    AgentMiddleware,
    AgentState,
    ModelRequest,
    ModelResponse,
)
from langgraph.runtime import Runtime
from typing import Any, Callable

class LoggingMiddleware(AgentMiddleware):
    def before_model(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
        print(f"About to call model with {len(state['messages'])} messages")
        return None

    def after_model(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
        print(f"Model returned: {state['messages'][-1].content}")
        return None

    async def abefore_model(
        self, state: AgentState, runtime: Runtime
    ) -> dict[str, Any] | None:
        # Async version of before_model
        return None

    async def aafter_model(
        self, state: AgentState, runtime: Runtime
    ) -> dict[str, Any] | None:
        # Async version of after_model
        print(f"Model returned: {state['messages'][-1].content}")
        return None


agent = create_agent(
    model="gpt-5.5",
    middleware=[LoggingMiddleware()],
    tools=[...],
)

클래스를 쓸 때: 같은 훅의 동기·비동기 구현이 모두 필요할 때, 하나의 미들웨어에 여러 훅이 필요할 때, 복잡한 설정(설정 가능한 임계값, 커스텀 모델 등)이 필요할 때, 초기화 시 설정으로 여러 프로젝트에 재사용할 때 좋아요.

커스텀 상태 스키마 (Custom state schema)

미들웨어가 훅 사이에서 상태를 추적해야 한다면, 에이전트 상태를 커스텀 속성으로 확장할 수 있어요. 이렇게 하면:

• 실행 전반에 걸쳐 카운터·플래그 같은 값을 유지하고
• before_model에서 after_model로 또는 다른 미들웨어 사이에서 데이터를 공유하며
• 요율 제한, 사용량 추적, 사용자 컨텍스트, 감사 로깅 같은 횡단 관심사를 핵심 에이전트 로직 수정 없이 추가하고
• 축적된 상태를 바탕으로 실행을 계속할지, 다른 노드로 점프할지, 동작을 동적으로 바꿀지 결정할 수 있어요.

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


class CustomState(AgentState):
    model_call_count: NotRequired[int]
    user_id: NotRequired[str]


@before_model(state_schema=CustomState, can_jump_to=["end"])
def check_call_limit(state: CustomState, runtime: Runtime) -> dict[str, Any] | None:
    count = state.get("model_call_count", 0)
    if count > 10:
        return {"jump_to": "end"}
    return None


@after_model(state_schema=CustomState)
def increment_counter(state: CustomState, runtime: Runtime) -> dict[str, Any] | None:
    return {"model_call_count": state.get("model_call_count", 0) + 1}


agent = create_agent(
    model="gpt-5.5",
    middleware=[check_call_limit, increment_counter],
    tools=[],
)

# Invoke with custom state
result = agent.invoke({
    "messages": [HumanMessage("Hello")],
    "model_call_count": 0,
    "user_id": "user-123",
})

실행 순서 (Execution order)

여러 미들웨어를 쓸 때 실행되는 방식을 이해해야 해요.

agent = create_agent(
    model="gpt-5.5",
    middleware=[middleware1, middleware2, middleware3],
    tools=[...],
)

핵심 규칙:

• before_* 훅: 앞에서 뒤로 (첫 번째부터 마지막까지)
• after_* 훅: 뒤에서 앞으로 (역순)
• wrap_* 훅: 중첩 (첫 번째 미들웨어가 나머지를 모두 감쌈)

에이전트 점프 (Agent jumps)

미들웨어에서 일찍 빠져나오려면 jump_to가 있는 딕셔너리를 반환하세요. 점프 가능한 대상:

• 'end' — 에이전트 실행의 끝으로 점프 (또는 첫 번째 after_agent 훅)
• 'tools' — 도구 노드로 점프
• 'model' — 모델 노드로 점프 (또는 첫 번째 before_model 훅)

from langchain.agents.middleware import after_model, hook_config, AgentState
from langchain.messages import AIMessage
from langgraph.runtime import Runtime
from typing import Any


@after_model
@hook_config(can_jump_to=["end"])
def check_for_blocked(state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
    last_message = state["messages"][-1]
    if "BLOCKED" in last_message.content:
        return {
            "messages": [AIMessage("I cannot respond to that request.")],
            "jump_to": "end"
        }
    return None

모범 사례 (Best practices)

• 미들웨어는 하나의 일만 잘하게 유지하세요.
• 에러를 우아하게 처리하세요 — 미들웨어 에러가 에이전트를 죽이지 않게요.
• 알맞은 훅 타입을 쓰세요: 순차 로직(로깅, 검증)엔 노드 스타일, 제어 흐름(재시도, 폴백, 캐싱)엔 랩 스타일.
• 커스텀 상태 속성은 명확히 문서화하세요.
• 통합 전에 미들웨어를 독립적으로 단위 테스트하세요.
• 실행 순서를 고려해서 중요한 미들웨어를 리스트 앞쪽에 두세요.
• 가능하면 내장 미들웨어를 쓰세요.

더 알아보기 (Learn more)

• 미들웨어 개요
• 내장 미들웨어 (Prebuilt middleware)
• Middleware API reference