[Amazon Bedrock][LangChain] チャット会話履歴をセッション毎に記憶・保存する方法

みなさん、こんにちは！
福岡オフィスの青柳です。

Amazon Bedrockをプログラムから利用する際に、AWSが提供するSDKを使う方法がありますが、他にも、生成AI向けフレームワークである「LangChain」を使う方法もあります。

今回は、LangChainが提供する機能を使って以下のようなことを実現したいと思います。

• チャットの会話履歴を記憶して、生成AIに文脈を踏まえた回答をさせる。
• 会話履歴をDynamoDBに保存して、Lambda関数を呼び出す度に記憶がリセットされないようにする。
• 会話履歴のセッション管理を行うことで、複数ユーザーの会話履歴を生成AIが区別できるようにする。

今回使用した環境

今回使用した環境は、原則としてブログ執筆時点 (2024/04/13) の最新で、以下の通りです。

また、AWSのリージョンは「バージニア北部 (us-east-1)」を使用します。 (東京リージョンでは利用できるモデルが限られるため)

準備

Lambda関数の作成を始める前に、いくつか準備をします。

Bedrockの「モデル」を利用可能にする

Bedrockマネジメントコンソールの「モデルアクセス」から、利用したい「モデル」へのアクセスを有効化しておきます。

Lambdaレイヤーの作成

Lambda関数の標準状態ではLangChainのPythonパッケージが含まれていないため、Lambdaレイヤーを使ってパッケージが利用できるようにします。

Lambdaレイヤーの作成と利用の方法は、下記ブログ記事を参考にしてください。

Lambda関数で使用するPythonと同じバージョンのPythonが使用できる環境で、LangChainのパッケージを含むzipファイルを作成します。

$ mkdir python
$ pip install -t ./python langchain
$ zip -r langchain.zip ./python

作成したzipファイルを使って、Lambdaレイヤーを作成します。

Lambda関数の作成と設定

Lambda関数を作成した後、以下の設定を行います。

• 上記で作成したLambdaレイヤーを使うように設定する
• タイムアウト時間を「30秒」程度に設定する (デフォルトの「3秒」だとLLMが応答を返す前にタイムアウトしてしまうため)

IAMロールの許可ポリシー設定

LangChain経由でBedrockを利用する時も、AWS純正のSDK (boto3) を使う場合と同様にIAMポリシーの設定が必要です。

以下の記述をIAMロールの許可ポリシーに追加します。

{
    "Effect": "Allow",
    "Action": "bedrock:InvokeModel",
    "Resource": "*"
}

Lambda関数コードの記述

ステップバイステップで、LangChainの各機能を確認しながら進めていきます。

AWS SDK (boto3) を使ってBedrockを利用したことがある方は、AWS SDK利用時とLangChain利用時の違いを意識しながら進めると良いでしょう。

ステップ 1： LangChainを使ってBecrock/Claude 3の呼び出しを行う

まずは、会話履歴の記憶などを行わない、最も基本的な生成AIチャットの機能を実装します。

Lambda関数コードの全体は以下のようになります：

from langchain_community.chat_models import BedrockChat
from langchain_core.messages import HumanMessage
import json


# 使用するLLMモデルを宣言
chat = BedrockChat(
    model_id="anthropic.claude-3-haiku-20240307-v1:0",
    model_kwargs={
        "max_tokens": 1000,
        "temperature": 0.5,
    },
)


# Lambdaハンドラー
def lambda_handler(event, context):
    # for debug
    print(f"Received event: {json.dumps(event, ensure_ascii=False)}")

    # イベントJSONから入力パラメーターを取得
    input_text = event.get("input_text")

    # メインの処理を呼び出す
    output_text = chat_conversation(input_text)

    # Lambda関数のレスポンスを返す
    return {
        "output_text": output_text,
    }


# メイン処理
def chat_conversation(input_text: str) -> str:
    # LLMに与えるプロンプトを設定
    messages = [
        HumanMessage(
            content=input_text
        ),
    ]

    # LLMにプロンプトを与えて、応答を含む結果を得る
    result = chat.invoke(messages)
    # for debug
    print(f"Result: {result}")

    # 得られた結果から、LLMが返した最新の応答テキストを抽出する
    output_text = result.content

    return output_text

ポイント別に解説していきます。

パッケージのインポート

from langchain_community.chat_models import BedrockChat
from langchain_core.messages import HumanMessage

1行目では、Amazon Bedrockに対応したLangChainのサブパッケージをインポートするように記述しています。

LangChainではBedrockの他にも「OpenAI」「Azure OpenAI」「Google VertexAI」など様々な生成AIプラットフォームに対応しています。 Bedrock以外のプラットフォームを利用する場合には、それに応じた記述をするようにします。

2行目のHumanMessageについては後で説明します。

使用するLLMモデルの宣言

chat = BedrockChat(
    model_id="anthropic.claude-3-haiku-20240307-v1:0",
    model_kwargs={
        "max_tokens": 1000,
        "temperature": 0.5,
    },
)

model_idには、使用するLLMモデルの種類とバージョンを示す文字列を指定します。 記述内容は、Bedrockマネジメントコンソールの「プロバイダー」から確認できます。(対象モデルを選択して「APIリクエスト」欄の内容を確認する)

model_kwargsには、モデルに与える各種パラメーターを指定します。

LLMに与えるプロンプトの設定

    messages = [
        HumanMessage(
            content=input_text
        ),
    ]

LLMに与えるプロンプトの設定方法は、LangChainの特徴の一つです。

AWS SDK、あるいは、各LLMが提供する標準のAPIを使ってプロンプトを記述する際、例えば以下のような記述をしていたと思います。

System: あなたは優秀なAIアシスタントです。・・・

Human: ◯◯について教えてください。

Assistant:

あるいは、LLMのモデルの種類やバージョンの違いによっては、異なる記述形式の場合もあります。

一方でLangChainでは、langchain_core.messagesパッケージに用意されているSystemMessageHumanMessageAIMessageなどのスキーマを使うことで、統一された構造的な記述を行うことができます。

LLMの呼び出し、結果の取得

    # LLMに会話リストを与えて、応答を含む結果を得る
    result = chat.invoke(messages)
    # for debug
    print(f"Result: {result}")

    # 得られた結果から、LLMが返した最新の応答テキストを抽出する
    output_text = result.content

LLMを呼び出す際はinvoke()メソッドを使用します。

応答は以下のような構造体となっています。回答内容を取得するにはcontentを参照します。

content='＜LLMからの回答内容＞'
response_metadata={
    'model_id': 'anthropic.claude-3-haiku-20240307-v1:0',
    'usage': {
        'prompt_tokens': 21,
        'completion_tokens': 250,
        'total_tokens': 271
    }
}
id='run-xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx-x'

実行してみる

それでは、作成したLambda関数を実行してみましょう。 実行する際に、入力イベントデータとして以下のJSONを指定します。

{
  "input_text": "日本一高い塔を教えてください。"
}

以下のように、LLMが回答した内容が戻り値として得られると思います。

{
  "output_text": "日本一高い塔は東京スカイツリーです。\n\n東京スカイツリーの主な情報は以下の通りです:\n\n・・・(以下省略)"
}

ここまで確認できましたら、次のステップに進みましょう。

ステップ 2： 「Chain」モジュールを使う方法に変更する

次のステップとして、機能面は変更せずに、LangChainの記述方法を少し変更してみます。

from langchain_community.chat_models import BedrockChat
from langchain_core.messages import HumanMessage
from langchain.chains import ConversationChain
import json


# 使用するLLMモデルを宣言
chat = BedrockChat(
    model_id="anthropic.claude-3-haiku-20240307-v1:0",
    model_kwargs={
        "max_tokens": 1000,
        "temperature": 0.5,
    },
)


# Lambdaハンドラー
def lambda_handler(event, context):
    # for debug
    print(f"Received event: {json.dumps(event, ensure_ascii=False)}")

    # イベントJSONから入力パラメーターを取得
    input_text = event.get("input_text")

    # メインの処理を呼び出す
    output_text = chat_conversation(input_text)

    # Lambda関数のレスポンスを返す
    return {
        "output_text": output_text,
    }


# メイン処理
def chat_conversation(input_text: str) -> str:
    # Chainモジュール (複数のモジュールを組み合わせて処理を行わせる)
    chain = ConversationChain(
        llm=chat,
        verbose=True,
    )

    # LLMに与えるプロンプトを設定
    messages = [
        HumanMessage(
            content=input_text
        ),
    ]

    # LLMにプロンプトを与えて、応答を含む結果を得る
    result = chain.invoke(messages)
    # for debug
    print(f"Result: {result}")

    # 得られた結果から、LLMが返した最新の応答テキストを抽出する
    output_text = result.get("response")

    return output_text

「Chain」モジュール

「ステップ 1」のコードとの違いは、LangChainの「Chain」モジュールを使用している点です。

Chainモジュールとは、LLMを使った単体の処理に留まらず、複数の処理を連係して行うために用意されているモジュールです。 例えば、LLM呼び出しの前処理として「LLMへの質問内容に危険性が無いかをチェックする」、後処理として「LLMの回答内容を要約する」などの処理を組み合わせることができます。

当ステップではChainに1つの処理「LLM呼び出し」のみを組み込んでいるため、結果的に「ステップ 1」と処理内容は変わりません。 この後のステップで、Chainに他の処理を組み込むことによって機能を追加していきます。

なお、「ステップ 1」でLLMに対してinvokeメソッドを呼び出した結果と、今回Chainモジュールに対してinvokeメソッドを呼び出した結果では、応答の形式が異なります。注意してください。

{
    'input': [
        HumanMessage(content='＜LLMへの質問内容＞')
    ],
    'history': '',
    'response': '＜LLMからの回答内容＞'
}

実行してみる

Lambda関数を更新した後、「ステップ 1」と同様の手順で動作確認を行います。

「ステップ 1」と同様の結果になることを確認してください。

ステップ 3： 「Memory」モジュールを使って会話履歴を記憶・保存できるようにする

「ステップ 1」や「ステップ 2」で「日本一高い塔」について質問した後に、再度Lanbda関数を呼び出して次の質問をしてみます。

{
  "input_text": "それは何メートルありますか？"
}

以下のような回答が返ってくると思います。

{
  "output_text": "申し訳ありませんが、具体的にどのものの大きさについて聞いているのかがわかりません。一般的な質問では、正確な回答をするのが難しいです。もし特定の物体の大きさを知りたい場合は、その物体について詳しく教えていただけますと、より具体的な回答ができると思います。"
}

これは、「ステップ 1」から「ステップ 2」までのコードでは、会話履歴を扱うことができないためです。

Lambda関数コードに、会話履歴を記憶する機能を追加します。

from langchain_community.chat_models import BedrockChat
from langchain_core.messages import HumanMessage
from langchain.chains import ConversationChain
from langchain.memory import ConversationBufferMemory
from langchain_community.chat_message_histories import DynamoDBChatMessageHistory
import json


# 使用するLLMモデルを宣言
chat = BedrockChat(
    model_id="anthropic.claude-3-haiku-20240307-v1:0",
    model_kwargs={
        "max_tokens": 1000,
        "temperature": 0.5,
    },
)


# Lambdaハンドラー
def lambda_handler(event, context):
    # for debug
    print(f"Received event: {json.dumps(event, ensure_ascii=False)}")

    # イベントJSONから入力パラメーターを取得
    input_text = event.get("input_text")

    # メインの処理を呼び出す
    output_text = chat_conversation(input_text)

    # Lambda関数のレスポンスを返す
    return {
        "output_text": output_text,
    }


# メイン処理
def chat_conversation(input_text: str) -> str:
    # Historyモジュール (Memoryの内容を外部記憶を使って永続化する)
    history = DynamoDBChatMessageHistory(
        table_name="ChatMessageHistory",
        session_id="abcd1234",
    )

    # Memoryモジュール (会話履歴を記憶する)
    memory = ConversationBufferMemory(
        chat_memory=history,
        return_messages=True,
    )

    # Chainモジュール (複数のモジュールを組み合わせて処理を行わせる)
    chain = ConversationChain(
        llm=chat,
        memory=memory,
        verbose=True,
    )

    # LLMに与えるプロンプトを設定
    messages = [
        HumanMessage(
            content=input_text
        ),
    ]

    # LLMにプロンプトを与えて、応答を含む結果を得る
    result = chain.invoke(messages)
    # for debug
    print(f"Result: {result}")

    # 得られた結果から、LLMが返した最新の応答テキストを抽出する
    output_text = result.get("response")

    return output_text

「Memory」モジュール

LangChainの「Memory」モジュールを使うことで、会話履歴を記憶させることができます。

    memory = ConversationBufferMemory(
        chat_memory=history,
        return_messages=True,
    )

「Chain」モジュールのmemoryパラメーターを追加設定します。

    chain = ConversationChain(
        llm=chat,
        memory=memory,
        verbose=True,
    )

これで、LLM呼び出しの処理に加えて「会話履歴の記憶」の機能が追加されます。

「History」モジュール

「Memory」モジュールは会話履歴をメモリ上に保持しているだけですので、プロセスが終了すると会話履歴が消えてしまいます。

Lambda関数の場合それでは都合が悪いため、永続的なデータストアを使って会話履歴を保存することにします。

LangChainでは、さまざまなデータストアに対応したモジュールが用意されています。 今回は、AWSのDynamoDBを使って会話履歴を保存するためにDynamoDBChatMessageHistoryを使用します。

    history = DynamoDBChatMessageHistory(
        table_name="ChatMessageHistory",
        session_id="abcd1234",
    )

パラメーターtable_nameで保存先のDynamoDBテーブル名を指定します。(テーブルはこの後作成します)

また、session_idにはとりあえずダミーの文字列を設定しておきます。

会話履歴を保存するDynamoDBテーブルの作成

DynamoDBChatMessageHistoryで使用するDynamoDBは、パーティションキーを以下のように設定する必要があります。

• キー名： SessionId
• タイプ： 文字列

その他の設定は、適宜行なってください。

IAMロールの許可ポリシー設定

Lambda関数がDynamoDBへアクセスできるように、IAMロールの許可ポリシーに権限を追加します。

必要なアクションについて、DynamoDBChatMessageHistoryのドキュメントには詳しい記述がありませんでしたが、ソースを確認したところ「getItem」「putItem」「deleteItem」の権限があれば良さそうでした。(将来的に変わる可能性がありますので、チェックしてください)

{
    "Effect": "Allow",
    "Action": [
        "dynamodb:getItem",
        "dynamodb:putItem",
        "dynamodb:deleteItem"
    ],
    "Resource": "arn:aws:dynamodb:<region-name>:<account-id>:table/<DynamoDBテーブル名>"
}

実行してみる

それでは、更新したLambda関数を実行してみましょう。

まず、1回目の実行では以下のイベントJSONを指定します。

{
  "input_text": "日本一高い塔を教えてください。"
}

「ステップ 1」「ステップ 2」と同様に、以下のような回答が返ってくると思います。

{
  "output_text": "日本一高い塔は東京スカイツリーです。\n\n東京スカイツリーの主な情報は以下の通りです:\n\n・・・(以下省略)"
}

次に、2回目の実行では、以下のイベントJSONを指定します。

{
  "input_text": "それは何メートルありますか？"
}

どうでしょう？ 「ステップ 1」「ステップ 2」とは異なる回答が返ってきましたか？

{
  "output_text": "東京スカイツリーの高さは634メートルです。\n\n東京スカイツリーは2012年に完成した日本一高い電波塔で、世界で2番目に高い自立式電波塔となっています。・・・(以下略)"
}

このように、文脈を踏まえた回答が返ってくれば、「会話履歴の記憶・保存」が正常に動作していることが確認できました。

(参考) DynamoDBには「会話履歴」がどのように保存されているか？

会話履歴がどのようにDynamoDBに保存されているのか、見てみましょう。

長いので一部を抜粋して記載しています：

{
  "SessionId": "abcd1234",
  "History": [
   {
    "data": {
     "content": ＜1回目の会話： LLMへの質問内容＞
    },
    "type": "human"
   },
   {
    "data": {
     "content": ＜1回目の会話： LLMからの回答内容＞
    },
    "type": "ai"
   },
   {
    "data": {
     "content": ＜2回目の会話： LLMへの質問内容＞
    },
    "type": "human"
   },
   {
    "data": {
     "content": ＜2回目の会話： LLMからの回答内容＞
    },
    "type": "ai"
   }
  ]
 }

ステップ 4： セッション毎に会話履歴を管理できるようにする

「ステップ 3」で作成した環境はLambda関数を呼び出す度に会話を記憶・保存してくれますが、このままでは永遠に会話履歴が増えていくばかりです。

自分専用の「パーソナル生成AIチャットボット」であればこれでも良いのですが、いろんなユーザーが利用する場合にはこれでは困ります。

そこで、「History」モジュールの「セッションID」を設定することにより、会話履歴のセッション管理を行えるようにします。

from langchain_community.chat_models import BedrockChat
from langchain_core.messages import HumanMessage
from langchain.chains import ConversationChain
from langchain.memory import ConversationBufferMemory
from langchain_community.chat_message_histories import DynamoDBChatMessageHistory
import json
import uuid


# 使用するLLMモデルを宣言
chat = BedrockChat(
    model_id="anthropic.claude-3-haiku-20240307-v1:0",
    model_kwargs={
        "max_tokens": 1000,
        "temperature": 0.5,
    },
)


# Lambdaハンドラー
def lambda_handler(event, context):
    # for debug
    print(f"Received event: {json.dumps(event, ensure_ascii=False)}")

    # イベントJSONから入力パラメーターを取得
    session_id = event.get("session_id")
    input_text = event.get("input_text")

    # セッションIDが未設定の場合、新規にセッションIDをセットする (UUIDを使ったランダムな文字列)
    if session_id is None:
        session_id = str(uuid.uuid4())

    # メインの処理を呼び出す
    output_text = chat_conversation(session_id, input_text)

    # Lambda関数のレスポンスを返す
    return {
        "session_id": session_id,
        "output_text": output_text,
    }


# メイン処理
def chat_conversation(session_id: str, input_text: str) -> str:
    # Historyモジュール (Memoryの内容を外部記憶を使って永続化する)
    history = DynamoDBChatMessageHistory(
        table_name="ChatMessageHistory",
        session_id=session_id,
    )

    # Memoryモジュール (会話履歴を記憶する)
    memory = ConversationBufferMemory(
        chat_memory=history,
        return_messages=True,
    )

    # Chainモジュール (複数のモジュールを組み合わせて処理を行わせる)
    chain = ConversationChain(
        llm=chat,
        memory=memory,
        verbose=True,
    )

    # LLMに与えるプロンプトを設定
    messages = [
        HumanMessage(
            content=input_text
        ),
    ]

    # LLMにプロンプトを与えて、応答を含む結果を得る
    result = chain.invoke(messages)
    # for debug
    print(f"Result: {result}")

    # 得られた結果から、LLMが返した最新の応答テキストを抽出する
    output_text = result.get("response")

    return output_text

「セッションID」パラメーターの追加

Lambda関数のイベントJSONから受け取るパラメーターにsession_idを追加します。

session_idが指定されていれば、それをセッションIDとして使用します。

もしsession_idが未設定の場合は、UUIDを使って新規にセッションIDを生成します。

    # イベントJSONから入力パラメーターを取得
    session_id = event.get("session_id")
    input_text = event.get("input_text")

    # セッションIDが未設定の場合、新規にセッションIDをセットする (UUIDを使ったランダムな文字列)
    if session_id is None:
        session_id = str(uuid.uuid4())

「History」モジュールのセッションIDを可変にする

「ステップ 3」では、「History」モジュールのsession_idパラメーターにダミーの固定文字列を与えていました。 (会話履歴が永遠に記録され続けるのは、これが原因でした)

このステップでは、「イベントJSONから受け取ったセッションID」または「新規に生成したセッションID」の値をsession_idパラメーターに与えるように変更します。

    history = DynamoDBChatMessageHistory(
        table_name="ChatMessageHistory",
        session_id=session_id,
    )

これで、セッション管理が行えるようになりました。

実行してみる

それでは、最終形となったLambda関数を実行してみましょう。

まず、最初はsession_idを指定せずに、メッセージ内容だけを指定して実行します。

{
  "input_text": "こんにちは、私の名前はケンシロウです。"
}

応答は以下のようになると思います。

{
  "session_id": "11111111-1111-1111-1111-111111111111",
  "output_text": "こんにちは、ケンシロウさん。私の名前はClaudeです。とてもよろこんでお会いできました。ケンシロウさんは日本人ですか? 日本語を話されるのはとても素晴らしいですね。・・・(以下略)"
}

回答メッセージ内容に加えてsession_idが返ってきました。 これが、この会話のために新規に生成されたセッションIDです。

では続けて、このセッションIDを指定して次の会話を実行します。

{
  "session_id": "11111111-1111-1111-1111-111111111111",
  "input_text": "私の名前を言ってみてください。"
}

以下のように、最初に話した「名前」を生成AIが覚えていて、回答を返してくれると思います。

{
  "session_id": "11111111-1111-1111-1111-111111111111",
  "output_text": "はい、わかりました。お名前は「ケンシロウ」さんですね。私はこれまでの会話の中で、ケンシロウさんが日本人で日本語を話されることを学びました。・・・(以下略)"
}

ユーザーを変えて実行してみる

では、次は「別のユーザーが会話を開始した体で」Lambda関数を呼び出してみましょう。

再度、session_idを指定せずに、メッセージ内容だけを指定して実行します。

{
  "input_text": "こんにちは、私の名前はジャギです。"
}

回答が返ってきます。

{
  "session_id": "22222222-2222-2222-2222-222222222222",
  "output_text": "こんにちは、ジャギさん。私の名前はClaudeです。とても嬉しいです。あなたの名前は素敵ですね。・・・(以下略)"
}

今度は、さきほどとは異なるsession_idが返ってくると思います。

続けて、このセッションIDを指定して次の会話を実行します。

{
  "session_id": "22222222-2222-2222-2222-222222222222",
  "input_text": "私の名前を言ってみてください。"
}

以下のように、1人目と2人目を混同せずに区別して回答を返してくれるはずです。

{
  "session_id": "22222222-2222-2222-2222-222222222222",
  "output_text": "はい、分かりました。あなたの名前はジャギさんですね。私は正確に覚えています。・・・(以下略)"
}

これで、会話履歴を記憶・保存しつつ、セッション管理も行ってくれるチャットLambda関数の動作が確認できました。

おわりに

DynamoDBに会話履歴をセッション単位で保存する仕組みで、会話履歴とセッション管理の機能を持ったチャットシステム (っぽいもの) が作れました。

もしこれをLangChainを使わずに実装しようとすると、結構おおがかりなプログラムコードになることが予想されます。 しかし、LangChainを使えばシンプルな記述でこのような機能を持ったプログラムが作成できるのです。

LangChainには他にもまださまざまな機能が用意されていますので、機会があればそれらを使ったサンプルもご紹介していきたいと思います。

参考