(LangChain) LangGraph Tutorial - 기억 주입하기

LangGraph에서는 채크포인터(CheckPointer) 기능을 통해 챗봇의 상태를 저장하고 추적할 수 있습니다. 이를 활용하면 대화의 맥락을 유지하거나 에러 복구, Human-in-the-Loop(HITL), 그리고 Time-Travel 등 다양한 시나리오에서 유용하게 활용할 수 있습니다.

앞서 만든 Agent에서는 대화의 맥락을 이어갈 수 있는 형태의코드가 아니었습니다.

# agent.py

# ~~ 중략 ~~

# 그래프 빌더 초기화
graph_builder = StateGraph(AgentState, config_schema=GraphConfig)

# 'chatbot' 노드 추가
graph_builder.add_node("chatbot", chatbot)
# 'tool_node' 노드 추가
graph_builder.add_node("tools", tool_node)

# 'chatbot'을 엔트리포인트로 설정
graph_builder.set_entry_point("chatbot")

graph_builder.add_conditional_edges(
    "chatbot",
    route_tools,
    {"tools": "tools", END: END},
)

# Any time a tool is called, we return to the chatbot to decide the next step
graph_builder.add_edge("tools", "chatbot")

# 그래프 컴파일
graph = graph_builder.compile()

(결과)

================================ Human Message =================================

Hi there! My name is K.
================================== Ai Message ==================================

Hi K! How can I assist you today?
================================ Human Message =================================

Remember my name?
================================== Ai Message ==================================

I don't have the capability to remember personal information or previous interactions,
so I don't know your name. However, you can share it with me if you'd like!

Implementation

1. 체크포인터 만들기

MemorySaver를 사용하여 인메모리 체크포인터를 사용합니다.

# agent.py

from langgraph.checkpoint.memory import MemorySaver

memory = MemorySaver()
graph = graph_builder.compile(checkpointer=memory)

• MemorySaver 객체 를 생성하고, checkpointer를 사용하여 그래프를 컴파일 및 상태를 저장합니다.

튜토리얼에서는 in-memory checkpointer 를 사용하지만,
프로덕션 단계에서는 SqliteSaver 또는 PostgresSaver 를 사용하여 자체 DB에 연결할 수 있습니다.

2. ToolNode와 tools_condition을 사용하여 그래프 구축하기

병렬 처리, 조건부 흐름 제어, 그래프 기반 동적 흐름 관리 등 시스템의 효율성, 확장성, 유연성을 향상시키고 다양한 동작을 효율적으로 처리할 수 있도록 하기 위해 prebuilt 된 ToolNode 와 tools_condition 으로 변경합니다.

• ToolNode : 그래프 상태(메시지 목록 포함)를 입력으로 받고 도구 호출의 결과로 상태 업데이트를 출력하는 LangChain Runnable
• tools_condition : state의 가장 최근 message를 가져와 tool_calls가 존재하는지 여부에 따라 분기

BasicToolNode → ToolNode로 변경

# nodes.py

import os, json
from langchain_openai import ChatOpenAI
from my_ai_agent.utils.state import AgentState
from my_ai_agent.utils.tools import tools
from dotenv import load_dotenv
from langchain_core.messages import ToolMessage
from langgraph.prebuilt import ToolNode, tools_condition

load_dotenv()

llm = ChatOpenAI(
    model="gpt-4o-mini",
    api_key= os.getenv("OPENAI_API_KEY"),
    max_tokens=None,
    temperature=0.7,
)

# LLM에 도구를 바인딩
llm_with_tools = llm.bind_tools(tools)

system_prompt = "Chat with the AI assistant. You can ask questions about anything else."

# 'chatbot' 노드
def chatbot(state):
    """상태에서 메시지를 받아서 LLM을 호출하는 함수"""
    messages = state["messages"]
    messages = [{"role": "system", "content": system_prompt}] + messages
    response = llm_with_tools.invoke(messages)
    return {"messages": [response]}


# 도구 노드 생성    
tool_node = ToolNode(tools)

여기서부터는 gpt3.5로 하면 정확성이 너무 떨어집니다. 그나마 비용이 저렴한 40-mini로 바꿔서 진행합니다.

route_table → tools_condition 함수로 변경

# agent.py

import os
from typing import Literal, TypedDict
from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import MemorySaver
from my_ai_agent.utils.state import AgentState
from my_ai_agent.utils.nodes import chatbot, tool_node
from langgraph.prebuilt import tools_condition


# Define the config
class GraphConfig(TypedDict):
    model_name: Literal["anthropic", "openai"]


# 그래프 빌더 초기화
graph_builder = StateGraph(AgentState, config_schema=GraphConfig)

# 'chatbot' 노드 추가
graph_builder.add_node("chatbot", chatbot)
# 'tool_node' 노드 추가
graph_builder.add_node("tools", tool_node)


# 'chatbot'을 엔트리포인트로 설정
graph_builder.set_entry_point("chatbot")


# 'chatbot'에서 'tool_node'로 이동하는 조건 추가
# The `tools_condition` function returns "tools" if the chatbot asks to use a tool, and "END" if
# it is fine directly responding. This conditional routing defines the main agent loop.
graph_builder.add_conditional_edges(
    "chatbot",
    tools_condition,
    # The following dictionary lets you tell the graph to interpret the condition's outputs as a specific node
    # It defaults to the identity function, but if you
    # want to use a node named something else apart from "tools",
    # You can update the value of the dictionary to something else
    # e.g., "tools": "my_tools"
    {"tools": "tools", END: END},
)

# Any time a tool is called, we return to the chatbot to decide the next step
graph_builder.add_edge("tools", "chatbot")


# 그래프 컴파일
memory = MemorySaver()
graph = graph_builder.compile(checkpointer=memory)

LangGraph Cloud나 LangGraph Studio를 사용하는 경우, 그래프를 컴파일할 때 자동으로 checkpointer을 포함한 지속성 및 상태 관리를 처리하기 때문에 선언할 필요는 없습니다.

3. 그래프 실행

이전 메시지를 기억해서 답변 하는지 확인합니다.

# agent.py

# 그래프 실행
if __name__ == "__main__":

    config = {"configurable": {"thread_id": "USER_01"}}
    user_input = {"type": "user", "content": "Hi there! My name is K."}
    for event in graph.stream({"messages": [user_input]}, config, stream_mode="values"):
        event["messages"][-1].pretty_print()

    user_input = {"type": "user", "content": "Remember my name?"}
    for chunk in graph.stream({"messages": [user_input]}, config, stream_mode="values"):
        chunk["messages"][-1].pretty_print()

(결과)

================================ Human Message =================================

Hi there! My name is K.
================================== Ai Message ==================================

Hello, K! How can I assist you today?
================================ Human Message =================================

Remember my name?
================================== Ai Message ==================================

Yes, I remember your name is K. How can I assist you further?

4. Langgraph studio에서 확인해보기

LangGraph Studio를 열면 자동으로 새 스레드 창이 열립니다. 기존 스레드가 열려 있는 경우 다음 단계에 따라 새 스레드를 만듭니다.
우측 상단의 + 버튼을 클릭하여 새로운 스레드를 생성할 수 있습니다.

Reference

• Adding Memory to the chatbot
• 그래프에 영속성 추가하는 방법
• Migrating from ConversationalChain
• Prebuilt components