Amazon DynamoDBを利用して、LangChain.jsでの会話履歴を保存する

AlexaやLex / LINEなどのチャットbotもそうですが、この手の対話式インターフェイスは原則としてステートレスに作られています。そのため「過去にどんな入力があって、どんな返答を返したか」はDynamoDBやCloudflare Workers KV / MomentoなどのDBまたはストレージに保存する必要があります。

今回はAmazon DynamoDBを会話履歴DBとして利用する方法を調べました。

DynamoDBのテーブルを作る

まずはテーブルを用意します。会話のセッションIDをキーにGET / PUTを行うイメージですので、設定自体はかなりシンプルにしました。

LangChain.jsからDynamoDBをよび出す

LangChainがDB呼び出しなども抽象化してくれます。が、内部的にAWS SDKが必要ですので、インストールしておきましょう。

npm i @aws-sdk/client-dynamodb

会話履歴の保存は、LangChainではmemoryとよびます。セッションIDを今回は暫定でDateから生成していますが、なにかしらキーにできるID（ユーザーのSIDなど）があればそちらを使う方が良さそうです。

import { BufferMemory } from "langchain/memory";
import { DynamoDBChatMessageHistory } from "langchain/stores/message/dynamodb";

const memory = new BufferMemory({
  chatHistory: new DynamoDBChatMessageHistory({
    tableName: "langchain",
    partitionKey: "id",
    sessionId: new Date().toISOString(), // Or some other unique identifier for the conversation
    config: {
      region: "us-east-2",
      credentials: {
        accessKeyId: "<your AWS access key id>",
        secretAccessKey: "<your AWS secret access key>",
      },
    },
  }),
});

あとはChainにmemoryとしてつないでやるだけです。

    const model = new BedrockChat({
        model: "anthropic.claude-v2",
        region: env.aws.region,
        credentials: {
          accessKeyId: env.aws.credentials.accessKeyId,
          secretAccessKey: env.aws.credentials.secretAccessKey,
        },
    });

    const conversationChain = new ConversationChain({ llm: model, memory });
    const response = await conversationChain.invoke({
        input: "Stripeでのサブスクリプションの作り方"
    });

どのKVS / DBを利用するかを抽象化できているので、あとから変更するのもかなり簡単そうです。

どのようにデータが保存されるか

先ほどのコードを動かして、どのようにデータがDynamoDBに送られるか見てみました。

messagesにList形式で履歴を入れている様子です。もし外から使う必要が出た場合や、Streamで分析・Fine Tuningのデータソースに使ったりしたい場合は、この構造を覚えておくとよいでしょう。

セッションを復元させてみる

前回の会話内容を元にやりとりができるようにしてみましょう。今回はsessionIdをDynamoDBに保存されているアイテムのIDに直接変更して実現させます。

const memory = new BufferMemory({
    chatHistory: new DynamoDBChatMessageHistory({
      tableName: env.aws.dynamodb.tableName,
      partitionKey: env.aws.dynamodb.partitionKey,
      sessionId: "2023-12-24T14:48:21.943Z", // Or some other unique identifier for the conversation
      config: {
        region: env.aws.region,
        credentials: {
            accessKeyId: env.aws.credentials.accessKeyId,
            secretAccessKey: env.aws.credentials.secretAccessKey,
        },
      },
    }),
  });

せっかくなので、コンテキストが必要そうな質問文に変えておきましょう。

    const conversationChain = new ConversationChain({ llm: model, memory });
    const response = await conversationChain.invoke({
        input: "JavaScriptで教えてください"
    });
    console.log(response);

実行結果がこちらです。「Stripe」と明言していないにも関わらず、Stripeに関する回答が出てきました。

{
  "response": " はい、JavaScriptを使ってStripeでサブスクリプションを実装する方法について説明します。\n\nまずStripeアカウントを取得し、公開"
}

DynamoDB側にも履歴が追加されています。

余談

Alexaの時の経験上、この手のデータストアはExpireさせる仕組みが必要です。仕組みを設けない限り、セッションの数だけデータが増えていくため、2〜3年後くらいに無料枠を使い切ることがあります。そしてその頃になると、どのデータを消すべきかの判断が難しくなるので、課金するかアプリごと消すかの判断を迫られます。

「揮発して欲しいデータ」であることを踏まえると、LangChain.jsでMemoryを使うならばMomentoも検討するのがよいかもしれません。ただしこちらはこちらでキャッシュの上限が24時間なので、それを超える期間で保存したい場合には、TTL相当の機能を持つデータストアを探すか、機構を作るかになりそうです。

参考

https://js.langchain.com/docs/integrations/chat_memory/dynamodb