Amazon BedrockとStreamlitで簡単チャットアプリを作成する方法

Bedrockの素敵なところは、LangChainを始めとするOSSのエコシステムのおかげてOpenAIからの置き換えが簡単にできるところだと思います。チャットアプリもOpenAI版をちょっと変えるだけで実現できました。

これまでOpenAIで検証してきた過去の資産を活かしたままBedrock化ができるって素敵

Streamlitによるチャット

こちらのサンプルコードをもとにします。

まず、OpenAI版のコードがこちらとなります。

from langchain.chains import LLMChain
from langchain.llms import OpenAI
from langchain.memory import ConversationBufferMemory
from langchain.memory.chat_message_histories import StreamlitChatMessageHistory
from langchain.prompts import PromptTemplate

import streamlit as st

# Optionally, specify your own session_state key for storing messages
msgs = StreamlitChatMessageHistory(key="chat_messages")

memory = ConversationBufferMemory(memory_key="history", chat_memory=msgs)
if len(msgs.messages) == 0:
    msgs.add_ai_message("How can I help you?")


template = """You are an AI chatbot having a conversation with a human.

{history}
Human: {human_input}
AI: """
prompt = PromptTemplate(input_variables=["history", "human_input"], template=template)

# Add the memory to an LLMChain as usual
llm_chain = LLMChain(llm=OpenAI(), prompt=prompt, memory=memory)

for msg in msgs.messages:
    st.chat_message(msg.type).write(msg.content)

if prompt := st.chat_input():
    st.chat_message("human").write(prompt)

    # As usual, new messages are added to StreamlitChatMessageHistory when the Chain is called.
    response = llm_chain.run(prompt)
    st.chat_message("ai").write(response)

これだけでチャットができる事自体がすごいですが、これをBedrock版に変えます。

Bedrock化

変更点はLLMsをOpenAIからBedrockに変更し、プロンプトを少し調整するだけです。

💡テンプレートは元のサンプルを活かしつつ、Amazon Bedrock Workshopを参考にしました。

  from langchain.chains import LLMChain
- from langchain.llms import OpenAI
+ from langchain.llms import Bedrock
  from langchain.memory import ConversationBufferMemory
  from langchain.memory.chat_message_histories import StreamlitChatMessageHistory
  from langchain.prompts import PromptTemplate

  import streamlit as st

  # Optionally, specify your own session_state key for storing messages
  msgs = StreamlitChatMessageHistory(key="chat_messages")

  memory = ConversationBufferMemory(memory_key="history", chat_memory=msgs)
  if len(msgs.messages) == 0:
      msgs.add_ai_message("How can I help you?")

- template = """You are an AI chatbot having a conversation with a human.
- 
- {history}
- Human: {human_input}
- AI: """
+ template = """You are an AI chatbot having a conversation with a human.
+ You answer always in japanese.
+
+ Current conversation:
+ <conversation_history>
+ {history}
+ </conversation_history>
+
+ Here is the human's next reply:
+ <human_reply>
+ {human_input}
+ </human_reply>
  """
  prompt = PromptTemplate(input_variables=["history", "human_input"], template=template)

  # Add the memory to an LLMChain as usual
- llm_chain = LLMChain(llm=OpenAI(), prompt=prompt, memory=memory)
+ llm_chain = LLMChain(llm=Bedrock(model_id='anthropic.claude-instant-v1'), prompt=prompt, memory=memory)

  for msg in msgs.messages:
      st.chat_message(msg.type).write(msg.content)

  if prompt := st.chat_input():
      st.chat_message("human").write(prompt)

      # As usual, new messages are added to StreamlitChatMessageHistory when the Chain is called.
      response = llm_chain.run(prompt)
      st.chat_message("ai").write(response)

超簡単

チャット履歴をDynamoDBに保持させる

上記サンプルでは、StreamlitChatMessageHistoryを使用しており、画面のリロードでチャット履歴が消えてしまいます。

チャット履歴をDynamoDBに永続化し、画面リロード時にも履歴が残るように対応します。

DynamoDBにテーブルを作成

パーティションキーをSessionIdとし、文字列型で定義します。
テーブル名は任意で指定します。（ここではlangchain-bedrock）

プログラムの修正

StreamlitChatMessageHistoryをDynamoDBChatMessageHistoryに変更します。

  from langchain.chains import LLMChain
  from langchain.llms import Bedrock
  from langchain.memory import ConversationBufferMemory
- from langchain.memory.chat_message_histories import StreamlitChatMessageHistory
+ from langchain.memory.chat_message_histories import DynamoDBChatMessageHistory
  from langchain.prompts import PromptTemplate

  import streamlit as st

  # Optionally, specify your own session_state key for storing messages
- msgs = StreamlitChatMessageHistory(key="chat_messages")
+ msgs = DynamoDBChatMessageHistory(table_name="langchain-bedrock", session_id="0")

  memory = ConversationBufferMemory(memory_key="history", chat_memory=msgs)
  if len(msgs.messages) == 0:
      msgs.add_ai_message("How can I help you?")


  template = """You are an AI chatbot having a conversation with a human.
  You answer always in japanese.

  Current conversation:
  <conversation_history>
  {history}
  </conversation_history>

  Here is the human's next reply:
  <human_reply>
  {human_input}
  </human_reply>

  """
  prompt = PromptTemplate(input_variables=["history", "human_input"], template=template)

  # Add the memory to an LLMChain as usual
  llm_chain = LLMChain(llm=Bedrock(model_id='anthropic.claude-instant-v1'), prompt=prompt, memory=memory)

  for msg in msgs.messages:
      st.chat_message(msg.type).write(msg.content)

  if prompt := st.chat_input():
      st.chat_message("human").write(prompt)

      # As usual, new messages are added to StreamlitChatMessageHistory when the Chain is called.
      response = llm_chain.run(prompt)
      st.chat_message("ai").write(response)

これだけ？！

履歴のセッションを分ける場合は、DynamoDBChatMessageHistoryのsession_idパラメーターを個別に指定するだけでOKです。

LINEボット化

Lambdaで動くようにカスタマイズし、LINEボットにするのも簡単です。

参考

Slackボット化

チャットではなくて画像生成ですが、こちらと今回のコードサンプルを組み合わせれば作成できると思います。