streamlitとlangchainでチャットAIを作る

はじめに

今さらですが、最近の学習でstreamlitを知りました、
爆速でChatGPTのようなサービスを作れるとのことで、その過程を記録していきます。

Streamlit
Pythonで簡単にデータアプリケーションを作成できるフレームワークです。データサイエンスや機械学習の結果を、簡単にインタラクティブなウェブアプリケーションとして表示するためのツールとして人気があります。公式ドキュメント

langchain
LangChainは、LLMを使ったアプリケーションを簡単に作成・統合するためのフレームワークです。 公式ドキュメント

成果物

依存リソース

pythonライブラリ

• streamlit
• langchain-aws
• langchain-community
• langchain

AWSリソース

• Cloud9
• DynamoDB

便宜上、Cloud9 を使用していますが、SageMaker Studio のコードエディタや Colab でも代替可能です。

SageMaker Studio使用する際に、下記の記事を参考にしてください。

また、履歴を保存するためのデータベースが必要です。今回は DynamoDB を使用しますが、他のデータベースでも基本的には代用可能です。

DynamoDB の設定は以下の通りです。

サンプルコードは下記の通りです。

index.py

import streamlit as st
from langchain_aws import ChatBedrock
from langchain_community.chat_message_histories import DynamoDBChatMessageHistory
from langchain_core.messages import HumanMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder

st.title("Bedrock チャット")

if "session_id" not in st.session_state:
    st.session_state.session_id = "session_id"

if "history" not in st.session_state:
    st.session_state.history = DynamoDBChatMessageHistory(
    table_name="ai-chat", session_id=st.session_state.session_id
    )

if "chain" not in st.session_state:
    prompt = ChatPromptTemplate.from_messages(
        [
          ("system", "あなたのタスクはユーザーの質問に明確に答えることです。"),
          MessagesPlaceholder(variable_name="messages"),
          MessagesPlaceholder(variable_name="human_message"),
        ]
    )
    
    # ChatBedrock
    chat = ChatBedrock(
      model_id="anthropic.claude-3-sonnet-20240229-v1:0",
      region_name="us-east-1",
      model_kwargs={"max_tokens": 1000},
      streaming=True,
    )
    
    chain = prompt | chat
    st.session_state.chain = chain

if st.button("履歴クリア"):
    st.session_state.history.clear()

for message in st.session_state.history.messages:
    with st.chat_message(message.type):
        st.markdown(message.content)

if prompt := st.chat_input("なんでも聞いて.."):
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        response = st.write_stream(
          st.session_state.chain.stream(
            {
                "messages": st.session_state.history.messages,
                "human_message": [HumanMessage(content=prompt)]
            },
            config={"configurable": {"session_id": st.session_state.session_id}},
          )
        )
    st.session_state.history.add_user_message(prompt)
    st.session_state.history.add_ai_message(response)

streamlit run index.py --server.port 8080.

参考文献