LangChain(JS)にて DynamoDB-Backed Chat Memoryを試す(2) トークン上限対策
はじめに
前回の記事
LangChain(JS)にて DynamoDB-Backed Chat Memoryを試す(1)
の続きです。
トークン上限に対応するために、別のMemoryを使えないか試したいと思います。
なお、前回はLangChainドキュメント見出しのDynamoDB-Backed Chat Memoryで機能名を表していましたが、今回の文中では他の機能と揃えて、クラス名のDynamoDBChatMessageHistoryと表記します。
BufferWindowMemory
公式ドキュメントのMemoryから該当する機能がないか探すと、BufferWindowMemoryが見つかりました。
前回使ったBufferMemory+DynamoDBChatMessageHistoryを置き換えてBufferWindowMemory+ DynamoDBChatMessageHistory にすると解決できそうです。
メモリークラスとユーティリティクラス
早速前回のコードをBufferWindowMemory置き換えてみようと思いますが、その前に１点補足しておきたい点があります。
BufferWindowMemoryとBufferMemoryはMemory機能の１つです。
しかし、今回（前回の記事でも）取り上げているDynamoDBChatMessageHistoryはこれらのMemory機能で使われる履歴データを操作（追加、一覧取得、削除など）するユーティリティクラスです。
名前だけ見るとBufferMemoryやBufferWindowMemoryと同レベルの機能のように見えますが、あくまで上記BufferMemory、BufferWindowMemoryの中で使われる１機能となります。
同じユーティリティクラスではChatMessageHistoryがあります。
これはBufferWindowMemoryとBufferMemoryにてデフォルトで使われています。
ちなみに、DynamoDBChatMessageHistoryとの違いは保存場所の違いです。（DynamoDBに読みっきするか、単に配列に読み書きをするかの違い）
大体想像がついてしまうかと思いますが、履歴を保存するストレージごとにユーティリティクラスが用意されているわけですね。
BufferMemory+DynamoDBChatMessageHistory(前回)
書き換える元となるコードは
こちら
になります。
BufferWindowMemory+ DynamoDBChatMessageHistory（今回）
単にBufferMemoryをBufferWindowMemoryに置換し、BufferWindowMemoryのコンストラクタに渡す引数にk: 10を加えています。
このkはLLMに渡す履歴の数（Human/AI、それぞれの発話一組で１とカウント)になります。
import { BufferWindowMemory } from "langchain/memory";
import { DynamoDBChatMessageHistory } from "langchain/stores/message/dynamodb";
import { ChatOpenAI } from "langchain/chat_models/openai";
import { ConversationChain } from "langchain/chains";
import dayjs from "dayjs";
import utc from "dayjs/plugin/utc.js";
import timezone from "dayjs/plugin/timezone.js";
import dotenv from "dotenv";
// OPEN AI のAPIキーなどを .envに定義しておく
dotenv.config();
// 現在の日付を取得
dayjs.extend(utc);
dayjs.extend(timezone);
dayjs.tz.setDefault("Asia/Tokyo"); // 日本のタイムゾーン
const currentDate = dayjs.tz().format("YYYY-MM-DD");
//const currentDate = dayjs.tz().format("YYYY-MM-DD HH:mm:ss");
const chatHistory = new DynamoDBChatMessageHistory({
tableName: "langchain",
partitionKey: "id",
sessionId: currentDate,
config: {
region: "ap-northeast-1",
},
});
const memory = new BufferWindowMemory({
k: 3,
chatHistory: chatHistory,
});
const model = new ChatOpenAI();
const chain = new ConversationChain({ llm: model, memory, verbose: true });
try {
const res1 = await chain.call({ input: "これまでのあなたとの会話をまとめてください。" });
console.log({ res1 });
} catch (error) {
if (error.response) {
console.log("Status:", error.response.status);
console.log("Data:", error.response.data);
} else if (error.request) {
console.log("Request:", error.request);
} else {
console.log("Error:", error.message);
}
}
// ```BufferWindowMemory``` を通さず直接```DynamoDBChatMessageHistory```で履歴を取得
const messages1 = await memory.chatHistory.getMessages();
console.log(messages1)
これで前回発生した、トークン数の上限エラーは発生しにくくなりました。
ただし、kの値を増やしたり、１回の発言が長くなれば依然としてトークン数の上限の超過は発生しますし、LLMの回答が途中で途切れやすくもなります。
あくまで履歴の「件数」を制限するだけである点に注意です。
なお、
const messages1 = await memory.chatHistory.getMessages();
console.log(messages1)
で出力されるログはkによる件数の絞り込みはありません。
実行後DynamoDBのレコードを直接みて見ても、データ自体は維持されます。同じsessionIdを使い続ける限り、DyanmoDBに格納される履歴は追加され続けます。
トークン数での制限を検討したいですね。
実装の比較
参考程度にBufferMemoryとBufferWindowMemoryの実装の比較をしてみます。
主な違いは引数の違い（kの有無）とloadMemoryVariablesの実装です。
それぞれのloadMemoryVariablesを比べてみましょう。
async loadMemoryVariables(_values) {
const messages = await this.chatHistory.getMessages();
if (this.returnMessages) {
const result = {
[this.memoryKey]: messages,
};
return result;
}
const result = {
[this.memoryKey]: getBufferString(messages, this.humanPrefix, this.aiPrefix),
};
return result;
}
async loadMemoryVariables(_values) {
const messages = await this.chatHistory.getMessages();
if (this.returnMessages) {
const result = {
[this.memoryKey]: messages.slice(-this.k * 2),
};
return result;
}
const result = {
[this.memoryKey]: getBufferString(messages.slice(-this.k * 2), this.humanPrefix, this.aiPrefix),
};
return result;
}
取得した履歴を全件返すか、一部を返すかの違いになっています。
トークン数を制限するカスタムメモリーの作成
上記を見ているとMemoryのloadMemoryVariablesを修正すれば、プロンプトに渡す履歴の、トークン数での制限が可能そうです。
入力トークン数＋履歴トークン数＋返信用トークン数を確保しておけば、かなりの確率でトークン上限の超過は防げるのではないかと思います。
※LLMが長すぎる回答を作成して途切れるケースもあるので、これでもまだ万全ではありません。これについてはまた別に記事を書きたいと思います。
公式ドキュメントのCreating your own memory class
という項目で独自クラスの作成例が示されているので、これを参考にCustomBufferWindowMemoryというのを作ってみようと思います。
トークン数の計算方法については、別途記事を作成しているのでそちらを参考にしてください。
OpenAI API(Node.js)で トークン数の送信前チェックについて検証(1)
事前に以下のパッケージをインストールしてください。
npm install js-tiktokne
以下が作成したクラスです。
import { BaseChatMemory, getBufferString } from "langchain/memory";
import { getEncoding } from "js-tiktoken";
export class CustomBufferWindowMemory extends BaseChatMemory {
constructor(fields) {
super({
returnMessages: fields?.returnMessages ?? false,
chatHistory: fields?.chatHistory,
inputKey: fields?.inputKey,
outputKey: fields?.outputKey,
});
Object.defineProperty(this, "humanPrefix", {
enumerable: true,
configurable: true,
writable: true,
value: "Human",
});
Object.defineProperty(this, "aiPrefix", {
enumerable: true,
configurable: true,
writable: true,
value: "AI",
});
Object.defineProperty(this, "memoryKey", {
enumerable: true,
configurable: true,
writable: true,
value: "history",
});
Object.defineProperty(this, "totalToken", {
enumerable: true,
configurable: true,
writable: true,
value: 2000,
});
this.humanPrefix = fields?.humanPrefix ?? this.humanPrefix;
this.aiPrefix = fields?.aiPrefix ?? this.aiPrefix;
this.memoryKey = fields?.memoryKey ?? this.memoryKey;
this.totalToken = fields?.totalToken ?? this.totalToken;
// js-tiktoken
this.enc = getEncoding("cl100k_base");
}
calcToken(text) {
const value = this.enc.encode(text);
return value.length;
}
get memoryKeys() {
return [this.memoryKey];
}
async loadMemoryVariables(_values) {
const originMessages = await this.chatHistory.getMessages();
let messages = [];
let totalToken = 0;
for (let i = 0; i < originMessages.length - 1; i++) {
if (i % 2 == 1) {
// ペアで処理（AI/Human片方のみの発話は用いない）するので
// i が奇数の時は処理をskipし偶数の時にまとめて処理する
continue;
}
const endIndex = originMessages.length - i;
const startIndex = endIndex - 2;
const tempMesg = originMessages.slice(startIndex, endIndex);
let str = getBufferString(tempMesg, this.humanPrefix, this.aiPrefix);
totalToken += this.calcToken(str);
if (totalToken > this.totalToken) {
break;
}
messages = tempMesg.concat(messages);
}
if (this.returnMessages) {
const result = {
[this.memoryKey]: messages,
};
return result;
}
const result = {
[this.memoryKey]: getBufferString(
messages,
this.humanPrefix,
this.aiPrefix
),
};
return result;
}
}
処理の概要としては、新しい履歴から２件ずつ（AI & Human)履歴を取得し、トークン数を計算し、合計トークン寸が指定値を超えていなければ次の２件の履歴を取得します。
トークン数の計算の際には
'Human: human発言\n' +'AI: ai発言41\n'
のように発話者を示す文字列も含めて計算していますが、ここは作成するシステムで使うプロンプト次第で変更を検討してください。
なお、プロンプトに渡す履歴だけでなく保存する履歴の上限も検討が必要かと思いますが、記事が長くなってきたので今回は省きたいと思います。
改修を検討される場合はloadMemoryVariablesのようにsaveContextを実装するか、DynamoDBChatMessageHistoryの処理を修正することになるかと思います。
またソースコードを検索していて見つけたのですがOpenAIAgentTokenBufferMemoryクラスのsaveContextも参考になるかと思います。
他に参考になりそうなMemory
最後に、公式ドキュメント読んでいて見つけた、DynamoDBChatMessageHistoryと組み合わせて活用できそうなMemory機能を紹介したいと思います。
ConversationSummaryMemory/ConversationSummaryBufferMemory
会話の履歴を直接埋め込むのではなく、要約を埋め込みます。
２つの違いはトークン数による制限の有無とあります。
Conversation summary memory
ConversationSummaryBufferMemory
VectorStoreRetrieverMemory
入力された文章から過去の会話履歴を検索し、類似度が高いものを利用します。
EntityMemory
ドキュメントやサンプルコードを見ているだけでは少しわかりにくかったので、利用されているプロンプトを読んでみました。その中で以下の部分がわかりやすかったです。
プロンプト（一部）
You are an AI assistant helping a human keep track of facts about relevant people, places, and concepts in their life.
上記翻訳
あなたは、人々や場所、概念に関する事実を追跡し、それらに関する要点を人間が記録できるよう支援するAIアシスタントです。
このMemoryをそのまま使うのも良いですが、指定したトピックに限定したプロンプトにするなど応用ができそうですね。
CombinedMemory
複数のMemoryの内容を同時に扱うことができるようになります。
例えば、直近２〜３個の履歴と、それより少し長めの記憶として要約をプロンプトに組み込むような使い方でしょうか。
How to use multiple memory classes in the same chain
まとめ
今回はDynamoDBChatMessageHistoryをBufferMemoryだけでなくBufferWindowMemoryと組み合わせて利用してみました。
また、Memoryクラスの拡張を行い、埋め込む履歴をトークン数で絞り込めるようにしました。
実際に動作確認までしていませんが、他にもDynamoDBChatMessageHistoryと組み合わせて活用できそうなMemory機能も確認することができました。
実はまだMemory関連では気になる点や試した実装があるのですDynamoDBChatMessageHistoryに関する記事から大分はみ出してしまうので、また別の機会に記事を作ることができればと思います。
Discussion