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
置き換えてみようと思いますが、その前に1点補足しておきたい点があります。
BufferWindowMemory
とBufferMemory
はMemory機能の1つです。
しかし、今回(前回の記事でも)取り上げているDynamoDBChatMessageHistory
はこれらのMemory機能で使われる履歴データを操作(追加、一覧取得、削除など)するユーティリティクラスです。
名前だけ見るとBufferMemory
やBufferWindowMemory
と同レベルの機能のように見えますが、あくまで上記BufferMemory
、BufferWindowMemory
の中で使われる1機能となります。
同じユーティリティクラスではChatMessageHistoryがあります。
これはBufferWindowMemory
とBufferMemory
にてデフォルトで使われています。
ちなみに、DynamoDBChatMessageHistory
との違いは保存場所の違いです。(DynamoDBに読みっきするか、単に配列に読み書きをするかの違い)
大体想像がついてしまうかと思いますが、履歴を保存するストレージごとにユーティリティクラスが用意されているわけですね。
BufferMemory
+DynamoDBChatMessageHistory
(前回)
書き換える元となるコードは
こちら
になります。
BufferWindowMemory
+ DynamoDBChatMessageHistory
(今回)
単にBufferMemory
をBufferWindowMemory
に置換し、BufferWindowMemory
のコンストラクタに渡す引数にk: 10
を加えています。
このk
はLLMに渡す履歴の数(Human/AI、それぞれの発話一組で1とカウント)になります。
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の値を増やしたり、1回の発言が長くなれば依然としてトークン数の上限の超過は発生しますし、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;
}
}
処理の概要としては、新しい履歴から2件ずつ(AI & Human)履歴を取得し、トークン数を計算し、合計トークン寸が指定値を超えていなければ次の2件の履歴を取得します。
トークン数の計算の際には
'Human: human発言\n' +'AI: ai発言41\n'
のように発話者を示す文字列も含めて計算していますが、ここは作成するシステムで使うプロンプト次第で変更を検討してください。
なお、プロンプトに渡す履歴だけでなく保存する履歴の上限も検討が必要かと思いますが、記事が長くなってきたので今回は省きたいと思います。
改修を検討される場合はloadMemoryVariables
のようにsaveContext
を実装するか、DynamoDBChatMessageHistory
の処理を修正することになるかと思います。
またソースコードを検索していて見つけたのですがOpenAIAgentTokenBufferMemory
クラスのsaveContext
も参考になるかと思います。
他に参考になりそうなMemory
最後に、公式ドキュメント読んでいて見つけた、DynamoDBChatMessageHistory
と組み合わせて活用できそうなMemory機能を紹介したいと思います。
ConversationSummaryMemory/ConversationSummaryBufferMemory
会話の履歴を直接埋め込むのではなく、要約を埋め込みます。
2つの違いはトークン数による制限の有無とあります。
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の内容を同時に扱うことができるようになります。
例えば、直近2〜3個の履歴と、それより少し長めの記憶として要約をプロンプトに組み込むような使い方でしょうか。
How to use multiple memory classes in the same chain
まとめ
今回はDynamoDBChatMessageHistory
をBufferMemory
だけでなくBufferWindowMemory
と組み合わせて利用してみました。
また、Memoryクラスの拡張を行い、埋め込む履歴をトークン数で絞り込めるようにしました。
実際に動作確認までしていませんが、他にもDynamoDBChatMessageHistory
と組み合わせて活用できそうなMemory機能も確認することができました。
実はまだMemory関連では気になる点や試した実装があるのですDynamoDBChatMessageHistory
に関する記事から大分はみ出してしまうので、また別の機会に記事を作ることができればと思います。
Discussion