Zenn
🔖

RAG機能付きチャットボットを作ろう-5_チャット表示のMarkdown化

2024/11/28に公開
Python
Streamlit
LangChain
生成 AI
chromadb
tech

TL;DR

前回の記事で、streamlitのsession_stateを使って、プロンプトや回答の保存を行いました。本稿では

  • チャット履歴をMarkdown形式で表示します。

実装イメージ

実装コード

主な変更点は

  • st.writeからst.markdownに変更

コード

import streamlit as st
from openai import OpenAI
from dotenv import load_dotenv
import os

load_dotenv(".env")

# OpenAIのAPIクライアントを初期化
# APIキーを環境変数から取得
client = OpenAI(
    api_key=os.environ['OPENAI_API_KEY']
    )

# プロンプトを入力すると、チャットボットが返答を返す関数を定義
# 入力はOpenAIのAPIクライアントとプロンプト
def default_chat(client, prompt):
    response = client.chat.completions.create(
        model="gpt-4o-mini", # 好きなモデルを選択
        messages=[
            {"role": "system", "content": "You are AI assistant."},
            {"role": "user", "content": prompt}
        ]
    )
    return response.choices[0].message.content

# streamlitのsession_stateにチャット履歴を保存する
# もしチャット履歴がなければ、空のリストを作成
if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []

# チャット履歴を表示する関数
def display_chat_history():
    for chat in st.session_state.chat_history:
        if chat["role"] == "user":
            st.markdown(
                # 背景をグレーにして、角を丸くする
                f'<div style="background-color: #f0f0f0; border-radius: 10px; padding: 10px;">'
                f"ユーザー: {chat['content']}"
                '</div>', unsafe_allow_html=True)

        else:
            st.markdown(
                # 背景を青にして、角を丸くする
                f'<div style="background-color: #cfe2ff; border-radius: 10px; padding: 10px;">'
                f"チャットボット: {chat['content']}"
                '</div>', unsafe_allow_html=True)

st.title('RAG機能付きチャットボットを作ろう')
st.write('streamlitを使ったUIの作成')

# チャット履歴を表示
display_chat_history()

prompt = st.text_area('プロンプト入力欄', )

button1, button2 = st.columns(2)
if button1.button('チャット'):
    chat_response = default_chat(client, prompt)
    # チャット履歴に追加
    # ユーザーの入力を追加、roleはuser
    st.session_state.chat_history.append({"role": "user", "content": prompt})
    # チャットボットの返答を追加、roleはsystem
    st.session_state.chat_history.append({"role": "system", "content": chat_response})
    st.rerun()
if button2.button('RAG'):
    st.write('RAGボタンがクリックされました')

上記をmain.pyとして保存し、streamlitを起動します。

streamlit run main.py

以下のような画面が表示されれば成功です。

alt text

解説

st.markdown

st.writeからst.markdownに変更しました。
st.markdownは、markdown記法を使ってテキストを表示する関数です。ChatGPTなどの生成AIでも生成した回答にmarkdown記法を使っていることがあるので、そのまま表示することができます。
以下のコードで、st.markdownを使って、チャット履歴を表示しています。

# チャット履歴を表示する関数
def display_chat_history():
    for chat in st.session_state.chat_history:
        if chat["role"] == "user":
            st.markdown(
                # 背景をグレーにして、角を丸くする
                f'<div style="background-color: #f0f0f0; border-radius: 10px; padding: 10px;">'
                f"ユーザー: {chat['content']}"
                '</div>', unsafe_allow_html=True)

        else:
            st.markdown(
                # 背景を青にして、角を丸くする
                f'<div style="background-color: #cfe2ff; border-radius: 10px; padding: 10px;">'
                f"チャットボット: {chat['content']}"
                '</div>', unsafe_allow_html=True)

unsafe_allow_html=Trueは、HTMLタグをそのまま表示するためのオプションです。st.markdownはデフォルトでHTMLタグをエスケープするので、unsafe_allow_html=Trueを指定することで、HTMLタグをそのまま表示することができます。

最後に、ユーザーの入力やチャットボットの返答をdivタグで囲み、背景色を変えています。ユーザーの入力はグレー、チャットボットの返答は青にしています。

リンク

参考

編集履歴

  • 2024/11/27: 初稿

Discussion