LangMemのQuick StartをDatabricks上でクイックに実行してみる

導入

LangChain社がAIエージェントの長期記憶を管理するためのライブラリ、LangMemをリリースしました。

以下、上記Blogより邦訳して抜粋。

本日、エージェントが長期記憶を通じて学習し、改善するのに役立つライブラリであるLangMem SDKをリリースします。

会話から情報を抽出し、迅速な更新を通じてエージェントの動作を最適化し、行動、事実、イベントに関する長期記憶を維持するためのツールを提供します。

そのコアAPIは、任意のストレージシステムおよび任意のAgentフレームワーク内で使用でき、LangGraphの長期メモリレイヤーとネイティブに統合されます。また、追加の長期記憶結果を無料で提供するマネージドサービスも開始していますので、本番環境での使用に興味がある方はこちらからご登録ください。

私たちの目標は、時間の経過とともによりスマートでパーソナライズされたAIエクスペリエンスを誰でも簡単に構築できるようにすることです。この作業は、ホストされたLangMemアルファサービスとLangGraphの永続的な長期記憶レイヤーの以前の作業に基づいています。

インストールするには、次のコマンドを実行するだけです。

pip install -U langmem

正直、ここだけ読んでもサッパリわからないので全文読むことをお薦めします。
※　ちなみに私は一度全文を読んだだけではよくわかりませんでした。。。

自分の理解だと、紹介文の通りエージェントの「長期記憶」を管理するためのライブラリであり、個々の（エージェントとの）会話記憶などを横断的に利用できるようにするためのものと認識しています。
例えば、LangGraphは既にCheckPointerというグラフ処理結果を記憶する仕組がありますが、これは「短期記憶」という位置づけであり、単一のエージェント（グラフ）処理の中でのみ利用される記憶領域です。つまり、他のLangGraphエージェント（グラフ）ではその記憶領域の内容を利用することができません（という理解をしています）。

LangMemは個々のエージェントとは独立した記憶管理を行い、異なるエージェント処理間でもその記憶領域を共通して利用できるものと理解しています。

※　正直、まだ十分に理解が及んでいないので誤った理解かもしれません。是非正しい内容を指摘ください。。。

個人的にこういった概念の理解は実践を通じて理解が深まると考えているので、実際にQuick Startを動かしてみます。

実行はDatabricks on AWS上でノートブックを作成して行いました。
公式の内容はOpenAI社やAnthropic社のAPIを利用する構造になっていますが、今回はクイックにDatabricks単独で実行できるようにコードを変更しています。
（ただし、Databricks 基盤モデルのpay-per-token利用に対応した環境が必要です）
また、クラスタはサーバレスを利用しています。

Quick Startその1: Hot Path Quickstart Guide

以下の内容を実行します。

このガイドでは、LangMemのmanage_memoryツールを使用して自身の長期メモリを積極的に管理するLangGraphエージェントを作成します。

ノートブックを作成して、必要なパッケージをインストール。

%pip install langmem
%pip install langgraph databricks-langchain
%pip install "mlflow-skinny[databricks]>=2.20.2"

%restart_python

また、処理結果を精緻に追うためにMLflow LangChain Tracingを有効化します。

import mlflow

mlflow.langchain.autolog()

では、クイックスタートのAgent処理定義を実装。

LLM/Embedding計算には、Databricksのpay-per-tokensで提供されている基盤モデルAPIを利用しました。

処理内容自体はLangGraphのcreate_react_agent関数を使って単純な対話エージェントを作っています。ポイントは、その際にToolとしてcreate_manage_memory_tool関数で作成したメモリツールを指定していることです。
これがLangMemの長期記憶を管理する仕組みになるようです。

from langgraph.checkpoint.memory import MemorySaver
from langgraph.prebuilt import create_react_agent
from langgraph.store.memory import InMemoryStore
from langgraph.utils.config import get_store
from langmem import (
    # エージェントがメモリを作成、更新、削除できるようにする
    create_manage_memory_tool,
)
from databricks_langchain import ChatDatabricks
from databricks_langchain import DatabricksEmbeddings

embeddings = DatabricksEmbeddings(endpoint="databricks-bge-large-en")
llm = ChatDatabricks(endpoint="databricks-meta-llama-3-3-70b-instruct")


def prompt(state):
    """LLMのメッセージを準備する。"""
    # 設定されたcontextvarからストアを取得する;
    store = get_store()  # `create_react_agent`に提供されたものと同じ
    memories = store.search(
        # エージェントに設定したのと同じ名前空間内で検索する
        ("memories",),
        query=state["messages"][-1].content,
    )
    system_msg = f"""You are a helpful assistant.

## Memories
<memories>
{memories}
</memories>
"""
    return [{"role": "system", "content": system_msg}, *state["messages"]]


store = InMemoryStore(
    index={  # 抽出されたメモリを保存する
        "dims": 1024,
        "embed": embeddings,
    }
)
checkpointer = MemorySaver()  # グラフの状態をチェックポイントする

agent = create_react_agent(
    llm,
    prompt=prompt,
    tools=[  # メモリツールを追加
        # エージェントは "manage_memory" を呼び出して
        # IDでメモリを作成、更新、削除できます
        # 名前空間はメモリにスコープを追加します。ユーザーごとにメモリをスコープするには、("memories", "{user_id}") を使用します:
        create_manage_memory_tool(namespace=("memories",)),
    ],
    # メモリはこの提供された BaseStore インスタンスに保存されます
    store=store,
    # 各ノードの実行が完了した後にグラフの "state" がチェックポイントされ、
    # チャット履歴と耐久実行を追跡します
    checkpointer=checkpointer,
)

では、実行。

config = {"configurable": {"thread_id": "thread-1"}}

# エージェントを使用します。エージェントはまだメモリを保存していないので、
# 私たちのことを知りません
response = agent.invoke(
    {"messages": [{"role": "user", "content": "Know which display mode I prefer?"}]},
    config=config,
)
print("Outputその1: まだ知識がない状態の回答")
print(response["messages"][-1].content)

# Darkモードが好きだと教えます
response = agent.invoke(
    {"messages": [{"role": "user", "content": "dark. Remember that."}]},
    # 同じ thread_id を持つ config を使用して会話 (thread-a) を続けます
    config=config,
)
print("Outputその2: 知識を教えたときの反応")
print(response["messages"][-1].content)

# 記憶から好きなモードを引き出せるか質問
response = agent.invoke(
    {
        "messages": [
            {
                "role": "user",
                "content": "Hey there. Do you remember me? What are my preferences?",
            }
        ]
    },
    config=config,
)
print("Outputその3: 教えた内容で回答するか？")
print(response["messages"][-1].content)

出力結果

Outputその1: まだ知識がない状態の回答
Yes, I know you prefer a specific display mode.
Outputその2: 知識を教えたときの反応
I've updated the memory to reflect that you prefer dark display mode. I'll keep that in mind for our conversation.
Outputその3: 教えた内容で回答するか？
I remember you. You prefer dark display mode.

エージェントの問い合わせはそれぞれ独立していますが、教えた内容を記憶して最終的に回答しています。

また、MLflow Tracingの結果を見ると以下のようにmanage_memoryツールを呼びだして、記憶領域に情報を記録していることが見て取れます。

ちなみに、異なるthread_idを指定してエージェントを実行した場合、新しい会話処理という扱いとなり、以前の記憶は使用してくれません。
(逆に言えば、同じthread_idを使う限りは蓄積した記憶情報を利用してくれるハズ)

# 新しいスレッド = 新しい会話！
new_config = {"configurable": {"thread_id": "thread-2"}}
# エージェントは manage_memories ツールを使用して明示的に保存したものだけを
# 思い出すことができます
response = agent.invoke(
    {"messages": [{"role": "user", "content": "Hey there. Do you remember me? What are my preferences?"}]},
    config=new_config,
)
print(response["messages"][-1].content)

出力

It seems that the memory associated with the id "892cd9a6-4987-4ddf-9449-a24771e12533" has been updated. However, since the content of the memory is null, I don't have any information about your preferences. Could you please tell me more about yourself or what you would like me to remember about you?

Quick Startその2: Background Quickstart Guide

もう一つのクイックスタートも実行してみます。

このガイドでは、create_memory_store_manager を使用してバックグラウンドでメモリを抽出して統合する方法を示します。エージェントは、メモリがバックグラウンドで処理されている間、通常どおり続行します。

という始まりのように、非同期でLangMemの記憶領域を更新する方法のようですね。

では、Basic Usageを一部変更して実行します。

こちらもLangGraph Functional APIを使って単純なチャット処理を実装・実行しているのですが、create_memory_store_manager関数を使って作成したmemory_managerを使って対話結果を保管しているところがポイントのようです。

from langchain.chat_models import init_chat_model
from langgraph.func import entrypoint
from langgraph.store.memory import InMemoryStore

from langmem import ReflectionExecutor, create_memory_store_manager
from databricks_langchain import ChatDatabricks
from databricks_langchain import DatabricksEmbeddings

embeddings = DatabricksEmbeddings(endpoint="databricks-bge-large-en")
llm = ChatDatabricks(endpoint="databricks-meta-llama-3-3-70b-instruct")

store = InMemoryStore( # インメモリストアを作成
    index={
        "dims": 1024,
        "embed": embeddings,
    }
)  

# 会話からメモリを抽出するためのメモリマネージャーRunnableを作成
memory_manager = create_memory_store_manager(
    llm,
    # メモリを "memories" 名前空間 (ディレクトリ) に保存
    namespace=("memories",),  # 名前空間を指定
)

@entrypoint(store=store)  # LangGraphワークフローを作成
async def chat(message: str):
    response = llm.invoke(message)

    # memory_managerは会話履歴からメモリを抽出
    # OpenAIのメッセージ形式で提供
    to_process = {"messages": [{"role": "user", "content": message}] + [response]}
    await memory_manager.ainvoke(to_process)  # メモリを抽出
    return response.content
# 通常通り会話を実行
response = await chat.ainvoke(
    "I like dogs. My dog's name is Fido.",
)
print(response)

出力結果

That's great! Fido is a classic name for a dog. What kind of dog is Fido? Is he a big fluffy breed or a small energetic one? What's his personality like? I'd love to hear more about your furry friend!

memory_managerに保管された情報はストアのsearchメソッドで確認できるようです。

print(store.search(("memories",)))

出力結果

[Item(namespace=['memories'], key='888dc50b-2c28-4768-a1ea-7fa3ea461767', value={'kind': 'Memory', 'content': {'content': 'The user likes dogs and has a dog named Fido.'}}, created_at...

まとめ

LangMemのクイックスタートをDatabricksを使ってクイックに実行してみました。

こういったライブラリをLangChain社がリリースするということは、エージェント構築において長期記憶の管理が重要だという認識をしているから、ということだと考えています。
個人的にまだ十分理解できているとは言い難いため、深く理解できるように努めていきます。