langchain version0.3におけるメモリ機能の備忘録

langchain version0.3におけるメモリ機能は、ver0.2から大きく変わっています。
以前、ver0.2について軽くまとめていましたが、再度まとめなおしてみることにしました。

参考

v0.2からv0.3のメモリ機能の移行についてまとめられたページ
https://python.langchain.com/docs/versions/migrating_memory/

大事なこと

ver0.2で利用した関数は、当面廃止されないのでそのまま使って大丈夫です

どう変わったか概要

langgraphのメモリ機能が優秀なので、全部そっちを使う形に変更となりました。

八種類あったメモリ関連の関数は、以下のように移行できます。

①　langgraphのメモリ機能を使う

• ConversationBufferMemory
  • 会話の履歴を保持する
  • 会話の履歴を全て保持するため、会話が長くなるとプロンプトが膨大になる。

②　langchainのtrim_messagesで過去の発言を削除したうえで、langgraphのメモリ機能を使う

• ConversationBufferWindowMemory

  • 直近k個の会話履歴のみ保持する
  • 古い会話が削除されるため、プロンプトが膨大になる問題が解決する

• ConversationTokenBufferMemory

  • ConversationBufferWindowMemoryのトークン版。直近のk会話ではなく、直近のnトークンを保持する。
    ⇒trim_messagesの引数設定を変えるだけで実現可能

③　langgrahでサマリ関数を実装して過去履歴をサマライズさせる

• ConversationSummaryMemory
  • 過去会話を要約して保持する
• ConversationSummaryBufferMemory
  • 過去会話の要約と、直近の会話（トークン数で指定）を保持する

④langgraphの長期記憶を使う

• Entity Memory
  • エンティティに関してのみメモリを保持する
• Conversation Knowledge Graph Memory
  • ナレッジグラフを作って会話についての情報を保持する
• VectorStore-Backed Memory
  • 過去会話をベクターストアに保持し、関連する会話のみ探索して用いる

この記事では、①~③について、具体的に実装がどう変わったかをメモっていきます。

具体的に実装がどう変わったか

前提

Azure OpenAI resourceを使います。openAIのAPIを使う場合は、チュートリアルをそのままなぞればOKのはず。

google colabを使います。モデルは、gpt-4o-mini。
大事な情報の設定からモデル作成までは以下。大事な情報はgoogle colabのシークレットで管理しています。

! pip install langchain-core langgraph>0.2.27
! pip install -qU langchain-openai
from langchain_openai import AzureChatOpenAI
from google.colab import userdata
import os

os.environ["AZURE_OPENAI_API_KEY"] = userdata.get("AZURE_OPENAI_API_KEY")

model = AzureChatOpenAI(
    model = "gpt-4o-mini",
    azure_endpoint=userdata.get("AZURE_OPENAI_ENDPOINT"),
    azure_deployment=userdata.get("AZURE_OPENAI_DEPLOYMENT_NAME"),
    openai_api_version=userdata.get("AZURE_OPENAI_API_VERSION"),
    openai_api_type="azure"
)

基本のモデルとのやりとりなど

まず、langgrapnのメモリ機能なんぞやからみていきます。

さっきつくったモデルに、ユーザープロンプトをなげてみます。

from langchain_core.messages import HumanMessage
result=model.invoke([HumanMessage(content="こんにちわ。良い天気だね")])
result.content

結果は以下となります。LLMからの返答を受け取れています。

こんにちは！はい、いい天気ですね。お出かけの予定はありますか？

履歴を読み込ませる場合は、HumanMessageとAIMessageを交互にリストにいれて、model.invokeに投げます。

from langchain_core.messages import AIMessage

result = model.invoke(
    [
        HumanMessage(content="私の名前はケイトです"),
        AIMessage(content="こんにちはケイト！本当にいい天気ですね。何か特別な予定はありますか？"),
        HumanMessage(content="私の名前は？"),
    ]
)

result.content

出力は以下のようになり、最初に名乗った名前が回答に反映されています。

あなたの名前はケイトですね！他に何かお話ししたいことがありますか？

ConversationBufferMemoryにあたるもの

langgraphを用いて、最もシンプルなメモリ機能であったConversationBufferMemoryを実装します。
つまり履歴を全部覚えておくというやつです。
前準備として以下のように、langgraphのGraphを定義しておきます。

from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import START, MessagesState,StateGraph

workflow = StateGraph(state_schema=MessagesState)

# モデル呼び出し用関数
def call_model(state: MessagesState):
  response = model.invoke(state["messages"])
  return {"messages":response}

# workflowでGraphを定義
workflow.add_edge(START, "model")
workflow.add_node("model",call_model)

# 定義したグラフを実行できる状態に落とし込む
memory=MemorySaver()
app = workflow.compile(checkpointer=memory)

設定したGraphを使ってみます。最初にconfigを定義するのですが、ここでthread_idを定義できるが超便利。複数の相手に対しての発言を別々のものとして保存できます。

config = {"configurable": {"thread_id": "abc123"}}
query="こんにちわ。僕はボブ"
input_messages = [HumanMessage(query)]
output = app.invoke({"messages":input_messages}, config)
output["messages"][-1].pretty_print()

出力は以下。

================================== Ai Message ==================================

こんにちは、ボブさん！お元気ですか？何かお手伝いできることがありますか？

メモリ機能が働いているか確認します。

query="私の名前は？"
input_messages = [HumanMessage(query)]
output = app.invoke({"messages":input_messages}, config)
output["messages"][-1].pretty_print()

出力からちゃんと、ボブであることを把握したまま会話が続きます。

================================== Ai Message ==================================

あなたの名前はボブさんです。何か他にお話ししたいことがありますか？

これを続ければ、ConversationBufferMemoryのように延々と過去の履歴を保持したまま会話を行うことができます。もちろん、入力トークンは膨れ上がります。

ConversationBufferWindowMemoryにあたるもの

langchainのtrim_messagesを使うと、古い発言を削除することができます。
trim_messagesは以下のように使います。

from langchain_core.messages import SystemMessage, trim_messages

trimmer = trim_messages(
    max_tokens=5,
    strategy="last",
    token_counter=len, #ここがlenだとメッセージ数でtrimしてくれる
    include_system=True,
    allow_partial=False,
    start_on="human",
)

messages = [
    SystemMessage(content="you're a good assistant"),
    HumanMessage(content="hi! I'm bob"),
    AIMessage(content="hi!"),
    HumanMessage(content="I like vanilla ice cream"),
    AIMessage(content="nice"),
    HumanMessage(content="whats 2 + 2"),
    AIMessage(content="4"),
    HumanMessage(content="thanks"),
    AIMessage(content="no problem!"),
    HumanMessage(content="having fun?"),
    AIMessage(content="yes!"),
]

trimmer.invoke(messages)

出力はこちら。SystemMessageはそのままですが、HumanMessageが「thanks」からはじまっており、それより前の出力はなくなっています。

[SystemMessage(content="you're a good assistant", additional_kwargs={}, response_metadata={}),
 HumanMessage(content='thanks', additional_kwargs={}, response_metadata={}),
 AIMessage(content='no problem!', additional_kwargs={}, response_metadata={}),
 HumanMessage(content='having fun?', additional_kwargs={}, response_metadata={}),
 AIMessage(content='yes!', additional_kwargs={}, response_metadata={})]

max_tokens=5と設定したので直近５つが残るかと思いましたが、４つしかのこりません。SystemMessageも含めて5つというカウントになっています。これは、include_system引数をTrueにしているためです。include_system引数をFalseにすると、SystemMessageはtrimされません。かといって、メッセージが5つ残るかといえばそんなことはなく、残るメッセージは変わらず4つとなります。

ということで、この機能を使えば、古い履歴を削除しながら、会話するというConversationBufferWindowMemoryのようなことができるわけです。

実際にtrim_messageを使って実装が以下になります。
一つ前の節で説明した、Graph定義のときにモデル呼び出し用関数としてcall_model関数を定義しました。あそこでモデルにプロンプトを投げているので、プロンプトを投げる直前に、trimmerで履歴を削除するだけです。ちなみにtrimmerは前節で定義したものを流用しています。

from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import START, MessagesState,StateGraph

workflow = StateGraph(state_schema=MessagesState)

# モデル呼び出し用関数
def call_model(state: MessagesState):
  #履歴削除
  trimmed_messages = trimmer.invoke(state["messages"]) 
  response = model.invoke(trimmed_messages)
  return {"messages":response}

# workflowでGraphを定義
workflow.add_edge(START, "model")
workflow.add_node("model",call_model)

# 定義したグラフを実行できる状態に落とし込む
memory=MemorySaver()
app = workflow.compile(checkpointer=memory)

実際に使ってみます。名前を聞いてみますが、うまくtrimされていれば、名乗り部分は削除されているので、LLMは名前をしらない状態で回答を返してきます。

config = {"configurable": {"thread_id": "abc567"}}
query = "What is my name?"
language = "English"
# ここでさきほど定義した長いmessagesに、名前を質問するプロンプトを追加して投げている。
input_messages = messages + [HumanMessage(query)]
output = app.invoke(
    {"messages": input_messages, "language": language},
    config,
)
output["messages"][-1].pretty_print()

出力は以下です。名前は知らないと発言できていますね。

================================== Ai Message ==================================

I don't know your name. What would you like me to call you?

ConversationTokenBufferMemoryにあたるもの

ConversationTokenBufferMemoryは、メッセージ数ではなくトークン数でカウントして履歴を削除するメモリ機能でした。
これも、さきほどのtrim_messages関数を用いて実現することができます。
token_counterをlenからmodelに変更するだけで、トークン数での履歴削除が実現できます。

from langchain_core.messages import SystemMessage, trim_messages

trimmer = trim_messages(
    max_tokens=30, # token数を指定
    strategy="last",
    token_counter=model, #使用モデルのトークン数でカウント
    include_system=True,
    allow_partial=False,
    start_on="human",
)

messages = [
    SystemMessage(content="you're a good assistant"),
    HumanMessage(content="hi! I'm bob"),
    AIMessage(content="hi!"),
    HumanMessage(content="I like vanilla ice cream"),
    AIMessage(content="nice"),
    HumanMessage(content="whats 2 + 2"),
    AIMessage(content="4"),
    HumanMessage(content="thanks"),
    AIMessage(content="no problem!"),
    HumanMessage(content="having fun?"),
    AIMessage(content="yes!"),
]

trimmer.invoke(messages)

出力は以下のようになります。

[SystemMessage(content="you're a good assistant", additional_kwargs={}, response_metadata={}),
 HumanMessage(content='having fun?', additional_kwargs={}, response_metadata={}),
 AIMessage(content='yes!', additional_kwargs={}, response_metadata={})]

モデル呼び出し用関数への組み込みは、ConversationBufferWindowMemoryの時と変わらないため割愛。

ConversationSummaryMemoryにあたるもの

これは、LLMを使ってサマライズさせる仕組みをlanggraphで実装する形になります。

以下に履歴をサマライズするタイプのチャットボットの例がのっているのでこれに沿って実装してみます。
https://langchain-ai.github.io/langgraph/how-tos/memory/add-summary-conversation-history/#setup

from typing import Literal
from langchain_core.messages import SystemMessage, RemoveMessage
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import MessagesState, StateGraph, START, END

memory = MemorySaver()  # 会話のデータを保存するためのメモリを初期化

# `messages`キー（MessagesStateクラスが持つ）に加えて、`summary`属性を追加します
class State(MessagesState):
    summary: str  # 会話の要約を保持するための属性

# モデルを呼び出すための処理を定義します
def call_model(state: State):
    # 要約が存在する場合、それをシステムメッセージとして追加します
    summary = state.get("summary", "")
    if summary:
        system_message = f"前回の会話の要約: {summary}"
        messages = [SystemMessage(content=system_message)] + state["messages"]
    else:
        messages = state["messages"]  # 要約がない場合はそのまま使用します
    response = model.invoke(messages)  # モデルを呼び出し、応答を取得
    # メッセージをリスト形式で返す（既存のリストに追加されるため）
    return {"messages": [response]}

# 会話を続けるか要約を行うかを判断するための処理
def should_continue(state: State) -> Literal["summarize_conversation", END]:
    """次に実行するノードを返す"""
    messages = state["messages"]
    # メッセージが6個以上の場合は会話を要約
    if len(messages) > 6:
        return "summarize_conversation"
    # そうでない場合は会話を終了
    return END

# 会話を要約する処理を定義
def summarize_conversation(state: State):
    # 会話を最初に要約
    summary = state.get("summary", "")
    if summary:
        # 既存の要約がある場合はそれを拡張するシステムプロンプトを使用
        summary_message = (
            f"これまでの会話の要約: {summary}\n\n"
            "上記の新しいメッセージを考慮して要約を拡張してください:"
        )
    else:
        summary_message = "上記の会話の要約を作成してください:"

    messages = state["messages"] + [HumanMessage(content=summary_message)]
    response = model.invoke(messages)
    # 表示したくないメッセージを削除（最新の2件以外のメッセージを削除）
    delete_messages = [RemoveMessage(id=m.id) for m in state["messages"][:-2]]
    print(response.content)
    return {"summary": response.content, "messages": delete_messages}

# 新しいワークフローグラフを定義
workflow = StateGraph(State)

# 会話ノードと要約ノードを追加
workflow.add_node("conversation", call_model)
workflow.add_node(summarize_conversation)

# 会話を開始点として設定
workflow.add_edge(START, "conversation")

# 条件付きの遷移を追加
workflow.add_conditional_edges(
    # スタートノードとして`conversation`を使用。
    # これは、`conversation`ノードが呼ばれた後に実行される遷移
    "conversation",
    # 次に実行するノードを決定する関数を指定
    should_continue,
)

# `summarize_conversation`からENDへの通常の遷移を追加。
# `summarize_conversation`が実行された後、会話が終了
workflow.add_edge("summarize_conversation", END)

# 最後にワークフローをコンパイル
app = workflow.compile(checkpointer=memory)

要約メッセージをわかりやすくするために、messagesの中身を日本語にしてから、appになげてみます。

messages = [
    SystemMessage(content="you're a good assistant"),
    HumanMessage(content="こんにちわ、僕は太郎"),
    AIMessage(content="どうも!"),
    HumanMessage(content="僕はたこやきが好きなんだ"),
    AIMessage(content="いいね"),
    HumanMessage(content="宿題教えてよ。2+2は？"),
    AIMessage(content="4"),
    HumanMessage(content="ありがとう"),
    AIMessage(content="どういたしまして"),
    HumanMessage(content="たのしい？"),
    AIMessage(content="もちろん"),
]
config = {"configurable": {"thread_id": "abc789"}}
query = "僕の名前は？"
language = "Japanese"
# ここでさきほど定義した長いmessagesに、名前を質問するプロンプトを追加して投げている。
input_messages = messages + [HumanMessage(query)]
output = app.invoke(
    {"messages": input_messages, "language": language},
    config,
)
output["messages"][-1].pretty_print()

出力は以下。

================================== Ai Message ==================================

あなたの名前は太郎ですね。

サマリ部分は以下のようになります。お、要約が日本語でされましたね。ConversationSummaryMemoryは英語でしかできませんでしたが、このやり方だと日本語もできるようです。何度かやっていると英語になる場合もあったので、サマライズ関数内のプロンプトをいじると安定するかもしれません。

太郎さんが自分の名前を名乗り、たこ焼きが好きだと話しました。その後、宿題の質問として2+2の答えを尋ね、4と答えられました。太郎さんは感謝の意を示し、楽しんでいるかどうかを尋ねると、楽しんでいると返答しました。

ConversationSummaryBufferMemoryは、うまくサマリの文字数と制御したり、trim_messagesと組み合わせればできそうです。
v0.2と比べると、実装が面倒ですが、自由度が高いのでやりたいことができそうな印象です。

まとめ

• langchainのv0.2で実装されていたメモリ機能のうち、短期記憶にあたる5つのメモリ機能について、v0.3でどのように実装されるかまとめました
• 履歴関連については、trim_messages関数をいじればいかようにもなる形で使いやすそう
• サマライズはグラフ組んだりしなければいけなくて面倒ですが、カスタマイズ性が高そうです

以上。