🎃

AI エージェント メッセージ履歴のトリミングと要約 (LangGraph)

に公開
Python
LangGraph
tech

エージェントは会話を繰り返したり、複数のツールを呼び出したりすると、メッセージ履歴がすぐに肥大化します。その結果、トークン数の増加によるコスト増加や、LLMの制約による性能低下といった問題に直面します。そこで重要になるのが、メッセージ履歴の管理です。この記事では、その代表的な方法である「トリミング」と「要約」を紹介します。

メッセージ履歴のトリミング

トリミングとは、メッセージ履歴の一部を削除し、必要な部分だけを残すことでトークン数を削減する方法です。主に「最初」または「最後」から順に履歴を削っていきます。

pre_model_hook は、エージェントがモデルを呼び出す直前にメッセージ履歴を加工するための仕組みです。これを利用することで、入力前に履歴をトリミングしたり要約したりできます。以下の例では、trim_messagesを使って履歴を短くしています。

graph = create_react_agent(
    model,
    tools,
    pre_model_hook=pre_model_hook,
    checkpointer=checkpointer,
)

メッセージ履歴のトリミングはlangchain_core.messages.utils.trim_messagesのメソッドを使用して行います。

def pre_model_hook(state):
    trimmed_messages = trim_messages(
        state["messages"],
        strategy="last",
        token_counter=count_tokens_approximately,
        max_tokens=100,
    )
    return {"llm_input_messages": trimmed_messages}

trim_messagesで重要な引数は、

  • messages: トリミングを行う対象のオブジェクト
  • strategy: メッセージ履歴の最初もしくは最後のどちらからキープするかを選択。Default last
  • max_tokens: キープする最大のトークン数
  • token_counter: トークン数を数えるための関数.

実際にコードを見ていきましょう

モデルとツールの定義
from langchain_core.tools import tool
from langchain_google_genai import ChatGoogleGenerativeAI

＃ モデルの定義
model = ChatGoogleGenerativeAI(model='gemini-2.5-flash', api_key=api_key)

# ツールの定義
@tool
def weather(location: str):
  """
  Get the weather in a given location
  """
  if any([city in location.lower() for city in ['london', 'lnd']]):
    return f"It's cloudy in London"
  elif any([city in location.lower() for city in ['paris', 'par']]):
    return f"It's raining in Paris"
  elif location.lower() == 'tokyo':
    return f"It's sunny in Tokyo"
  else:
    return f"I don't know the weather in {location}"

トリミングの設定：上記の通りメッセージ履歴の編集を行う関数を渡します。

from langchain_core.messages.utils import (
    trim_messages,
    count_tokens_approximately
)

def pre_model_hook(state):
    trimmed_messages = trim_messages(
        state["messages"],
        token_counter=count_tokens_approximately,
        max_tokens=50,
    )
    print(f"trimmed_messages: {trimmed_messages}")
    
    return {"llm_input_messages": trimmed_messages}

ここでは、わかりやすくするために、max_tokensをかなり小さくしています。token_counterには、count_tokens_approximatelyを使います。それを、stateのllm_input_messagesに返します。これは、メッセージ履歴を上書き保存せず、プロンプトを呼び出すときに使用する一時的なメッセージ履歴です。

これをエージェントに渡していきます。

from langgraph.prebuilt import create_react_agent
from langgraph.checkpoint.memory import InMemorySaver

# 記憶を管理するための
checkpointer = InMemorySaver()
agent = create_react_agent(
    model=model,
    tools=[weather],
    pre_model_hook=pre_model_hook,
    checkpointer=checkpointer
)

グラフを描写してみます。

from IPython.display import display, Image
display(Image(agent.get_graph().draw_mermaid_png()))


LangGraphにおけるpre_model_hookの実行フロー

図のように、エージェントが呼び出される前に、pre_model_hookを呼び出して、モデルに渡すメッセージ履歴のトリミングを行います。

実際に呼び出してみます。

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

inputs = {"messages": [("user", "What's the weather in London?")]}
result = agent.invoke(inputs, config=config)

inputs = {"messages": [("user", "What's it known for?")]}
result = agent.invoke(inputs, config=config)

for message in result['messages']:
  message.pretty_print()

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

What's the weather in London?
================================== Ai Message ==================================
Tool Calls:
  weather (a10e5a9f-e015-4288-a966-3d47da31b522)
 Call ID: a10e5a9f-e015-4288-a966-3d47da31b522
  Args:
    location: London
================================= Tool Message =================================
Name: weather

It's cloudy in London
================================== Ai Message ==================================

It's cloudy in London
================================ Human Message =================================

What's it known for?
================================== Ai Message ==================================

Please tell me what "it" refers to. I need more information to answer your question.

出力されたメッセージを見てみると、itが何を指しているのか、わからないと返答しています。これは、モデルに渡すプロンプトが最大50トークン(max_tokens=50)と小さく設定しているため、プロンプトにロンドンという情報が含まれず、2回目の呼び出しが行われてしまいました。その結果、このエラーが出ています。max_tokensを使用用途に合わせて、適宜調節してください。

モデルによっては、メッセージ履歴が必ず human（ユーザー）から始まっていないと正しく処理できないことがあります。そのため、トリミングの際に順序が崩れるとエラーが発生します。これを避けるには、start_on="human" を指定して「必ずユーザーから始まるように」調整します。

このようなメッセージ履歴に対してトリミングを行い、

AI -> human -> AI の順になるとgeminiでは、エラーが出てしまいます。これを防ぐために、start_on という引数にhumanと指定して、トリミング後の形を調整します。

またend_onの引数で最後のメッセージタイプを指定することも同様にできます。

トリミングはシンプルですが、削除された部分に重要な情報が含まれていると会話が破綻することがあります。そこで有効なのが要約です。要約は、履歴を単に削除するのではなく、要点を残した短縮版を作る仕組みです。

メッセージ履歴の要約

先程使用した、pre_model_hookに要約のノードを追加します。

from langmem.short_term import SummarizationNode, RunningSummary
from langchain_core.messages.utils import count_tokens_approximately
from langgraph.prebuilt.chat_agent_executor import AgentState

summarization_node = SummarizationNode( 
    token_counter=count_tokens_approximately,
    model=model,
    max_tokens=384,
    max_summary_tokens=128,
    output_messages_key="llm_input_messages",
)
  • max_tokens: 要約するメッセージ履歴の範囲をトークン数で指定
  • token_counter: トークン数を数えるための関数
  • max_summary_tokens: 要約したときの出力トークン数
  • model: 要約するためのモデル

output_messages_keyは先程と同じllm_input_messagesを指定して、モデルを呼び出すためのメッセージプロンプトに保存します。

追加でAgentStateを使用します。
AgentState は要約結果を一時的に保存する仕組みです。これを使うことで、毎回履歴全体を要約する必要がなくなり、新しい情報だけを追加要約できます。その結果、処理の効率が大幅に向上します。

class State(AgentState):
    context: dict[str, RunningSummary]

これらを渡してエージェントを作成します。

agent = create_react_agent(
    model=model,
    tools=[weather],
    pre_model_hook=summarization_node, 
    state_schema=State, 
    checkpointer=checkpointer,
)

では実際に、エージェント呼び出して出力と現在の要約を確認してみます。contextキーで現在使用している要約を確認できます

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

inputs = {"messages": [("user", "I'm Alex,a high school student, living in LA")]}
response = agent.invoke(inputs, config=config)

inputs = {"messages": [("user", "I love surfing and swimming")]}
response = agent.invoke(inputs, config=config)

inputs = {"messages": [("user", "I'm moving to NYC in Spring for my college")]}
response = agent.invoke(inputs, config=config)

inputs = {"messages": [("user", "Where am I living now?")]}
response = agent.invoke(inputs, config=config)

# 最後の出力を確認
response["messages"][-1].pretty_print()
# モデルに渡した要約を確認
print("\nSummary:", response["context"]["running_summary"].summary)

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

You are currently living in LA.

Summary: This conversation involves **Alex**, who initially expressed a love for surfing and swimming.

Alex then announced that they are **moving from LA to NYC in the Spring for college**.

The AI congratulated Alex on this news and, acknowledging their interests, provided information about **surfing and swimming opportunities in NYC** (like Rockaway Beach and pools), noting it would be a change from LA.

サンプルのメッセージ履歴は短いものですがしっかりと必要な情報が集約され、エージェントは要約された情報をしっかり活用できています。

結論

今回紹介したトリミングと要約を活用することで、無駄なトークンを減らし、効率的にエージェントを動かすことができます。会話の文脈を維持しながらコストや制約を回避するために、ぜひ実装に取り入れてみてください。

参考文献

trim_messages - LangGraph
SummarizationNode - LangGraph

Discussion