LangGraph Agent에 장기기억(LTM)추가하기

장기기억이 추가된 Agent의 전체 그래프

LangGraph로 만든 Agent에 장기 기억을 추가하기 위해서는 기존 그래프에 새로운 노드와 도구를 추가해야한다.

• 노드
  • 대화 전 저장된 기억을 불러오는 노드 (load_memories)
• 도구
  • 관련 기억을 검색해서 가져오는 도구 (SearchRecallMemoriesTool)
  • 기억하는 도구 (SaveMemoriesTool)

load_memories 노드가 하는 일?

load_memories 노드부터 살펴보겠다. load_memories는 Agent가 작업을 처리하기 시작할때 이전 대화 내용을 기반으로 관련된 기억을 불러오는 노드다. 불러온 기억은 State에 [”recall_memories”]로 저장되고, 이는 나중에 기억과 관련된 SystemPrompt에 Context로서 제공된다.

load_memories를 어떻게 구현하고, 통합하나요?

def load_memories(state: StateSchema, config: RunnableConfig) -> StateSchema:
    convo_str = get_buffer_string(_get_state_value(state, "messages"))
    convo_str = tokenizer.decode(tokenizer.encode(convo_str)[:2048])
    recall_memories = search_recall_memories_tool.invoke(
        convo_str,
        config
    )
    return {
        "recall_memories": recall_memories,
    }

이전 대화에 대한 너무 많은 맥락이 한꺼번에 들어가지 않도록 적절한 토큰길이로 잘라 관련 기억을 불러오게 되는데, 기억을 불러오는데는 SearchRecallMemoriesTool을 그대로 이용한다. 이 함수가 불리는 시점 이후로 Graph의 State에는 messages말고도 recall_memories가 추가로 들어가게 된다.

{
    ...
    "recall_memories": ["2025-04-03: 사용자는 초밥을 좋아합니다"]
}

물론 이대로만 하면 문제가 생길 것이다. StateSchema 는 기본적으로 AgentState를 사용하고 있는데, AgentState에는 “recall_memories” 멤버 변수가 없으므로 아무 기억도 저장하지 못한다.

class AgentState(TypedDict):
    """The state of the agent."""

    messages: Annotated[Sequence[BaseMessage], add_messages]

    is_last_step: IsLastStep

    remaining_steps: RemainingSteps

AgentState를 상속받아 recall_memories를 정의해서 Graph의 StateSchema로 지정해주어야한다.

class LTMAgentState(AgentState):
    recall_memories: Sequence[str] = []

이제 State에서 recall_memories를 접근할 수 있으므로, 이를 Context로 제공할 수 있다. "사용자와 관련된 기억은 이러이러한것이니까~ 이걸 참고해서 대답해"라는 SystemPrompt를 작성할 수 있다.

def _get_prompt_runnable(prompt: SystemMessage) -> Runnable:
    def _get_messages_with_memories(state):
        messages = _get_state_value(state, "messages")
        recall_memories = _get_state_value(state, "recall_memories", "No relevant memories.")

        memory_message = SystemMessage(
            content="Recalled Memories: {recall_memories}".format(
                recall_memories="\n".join(recall_memories)
            )
        )
        return [memory_message] + messages

    prompt_runnable = RunnableCallable(
        lambda state: [prompt] + _get_messages_with_memories(state),
        name=PROMPT_RUNNABLE_NAME,
    )
    return prompt_runnable

예시 코드에서는 이해를 위해 아주 간단하게 “Recalled Memories: {}”로 들어갔지만 이 기억을 어떻게 다룰것인지에 대한 상세한 설명은 각자 보강하면된다. 기존의 SystemPrompt에 추가로 기억과 관련된 SystemPrompt가 추가되는 형태다.

graph.add_node("load_memories", load_memories)
graph.add_edge("load_memories", "agent")
graph.set_entry_point("load_memories")

load_memories가 준비되었으니 그래프에 통합한다.

 

다음은 도구들이다.

 

기억을 가져오는 도구

Search Recall Memories Tool

앞서 load_memories 노드에서 필요했던 것처럼…연관된 기억을 가져오는 도구가 필요하다. 그것을 우리는 SearchRecallMemoriesTool이라고 하겠다.

 

기억은 VectorDB의 Similarity Search를 이용해 가져오는데, 이때 특정 유저에 대한 기억만 가져오도록 필터링을 했다.

def get_user_id(config: RunnableConfig) -> str:
    user_id = config["configurable"].get("user_id")
    if user_id is None:
        raise ValueError("User ID not found in configuration.")
    return user_id

def search_recall_memories(user_id: str, query: str) -> list[str]:
    index = MemoryIndex()
    return index.search(user_id=user_id, query=query, k=6)

class SearchRecallMemoriesTool(BaseTool):
    name: ClassVar[str] = "search_recall_memories"
    description: ClassVar[str] = """
    Retrieve previously stored information about the user.
    """

    def _run(self, query: str, config: Annotated[RunnableConfig, InjectedToolArg]) -> list[str]:
        user_id = get_user_id(config)
        return search_recall_memories(user_id=user_id, query=query)

구현을 살펴보면 RunnableConfig에서 user_id를 찾아서 해당 user_id와 query로 관련된 기억을 6개정도 들고오는게 끝이다.

다음 도구로 넘어가기전에, MemoryIndex를 살펴보겠다.

class MemoryIndex():

    def __init__(self):
        self.client = chromadb.HttpClient(CHROMADB_URI)
        self.chroma = Chroma(
            # ...
        )

    def search(self, user_id: str, query: str, k: int) -> list[str]:
        documents = self.chroma.similarity_search(
            query, 
            k,
            filter={
                "user_id": user_id
            }
        )
        # 최신순으로 정렬
        documents = sorted(documents, key=lambda x: x.metadata["created_at"], reverse=True)
        return [f"{doc.metadata["created_at"]}: {doc.page_content}" for doc in documents]

    def save(self, user_id: str, memory: str):
        document = Document(
            page_content=memory,
            metadata={
                "user_id": user_id,
                "created_at": datetime.now(pytz.timezone("Asia/Seoul")).strftime("%Y-%m-%d")
            }
        )
        self.chroma.add_documents([document])

특별히 추가한게 있다면 기억한 날짜다. 사용자의 요구가 시시때때로 바뀔수 있기때문에 어떤 요구가 더 최신일지에 대한 추가정보를 제공해준다면 더 잘 답변할 수 있을거라고 판단했기때문이다.

이제부터 나를 두목님이라고 불러!! -> "사용자는 두목님이라고 불리길 원합니다."
아냐 이제부터 행님이라고 불러그냥. -> "사용자는 행님이라고 불리길 원합니다."

 

기억을 저장하는 도구

Save Memories Tool

기억을 저장하는 SaveMemoriesTool은 LLM이 저장해야겠다고 판단한 기억들을 VectorDB에 쌓아주면된다.

def save_memories(user_id: str, memory: str) -> str:
    index = MemoryIndex()
    index.save(user_id, memory)
    return memory

class SaveMemoriesTool(BaseTool):
    name: ClassVar[str] = "save_memories"
    description: ClassVar[str] = """
    Store important user information for future retrieval.
    """

    def _run(self, memory: str, config: Annotated[RunnableConfig, InjectedToolArg]) -> str:
        user_id = get_user_id(config)
        return save_memories(user_id, memory)

이들을 모두 통합하고, Agent에게 작업을 호출할때 user_id를 제대로 설정해주기만 한다면 유저별 장기기억기능을 간단히 붙여볼 수 있다.

config = {
    "configurable": {
        "thread_id": thread_id,
        "user_id": user_id,
    }
}

기억을 추가하면 어떻게 되나요?

대화를 하면서 여러가지 것들을 기억해달라고 명시적으로 얘기할 수도 있고, 뭔가를 좋아한다거나 싫어한다거나, 나의 직업같은 특징들을 알아서 저장해준다. (물론 이건 시스템 프롬프트를 어떻게 작성하냐에 따라 다르다)

이런 기억이 있을 경우에

장기기억을 추가할 경우, 이렇게 이전 맥락이 전혀없는 새로운 대화에서도 내가 두목이라고 부르라고 했던 것을 기억하는 것같은 모습을 보인다.