LlamaIndex 學習筆記 - Streamlit 做 RAG chatbot UI

嘗試的原始碼在 https://github.com/unclefomotw/llamaindex-try/tree/main/src/rag_bot_1

之前我的範例，問問題都是直接改程式碼重跑一遍。但用了 Streamlit 以後，能非常快速兜出一個 GUI，生出跟 ChatGPT 一樣的聊天網頁去對話，很漂亮！

Chatbot UI as a webpage made by Streamlit — 用 Streamlit 做的 RAG 聊天介面

當我說很快，是只要寫 20 幾行程式碼就好，而且很好讀！

環境 #

安裝 Streamlit 用 pip 就可了。這次執行是在本機上，沒有放到雲端

基本元件 #

詳見我的範例程式碼，最基本用下面 python 畫出聊天介面

import streamlit as st

if prompt := st.chat_input("任何魔方的問題都可以問我唷"):
    # st.session_state.messages.append({"role": "user", "content": prompt})

    with st.chat_message("user"):
        st.markdown(prompt)

    query_engine = st.session_state.query_engine
    with st.chat_message("assistant"):
        response = query_engine.query(prompt)
        st.write(response.response)
    # st.session_state.messages.append({"role": "assistant", "content": response.response})

st.chat_input() 跟 st.chat_message() 都是 UI 元件，分別代表下方的聊天框與上方的聊天訊息

st.chat_input() 會回傳使用者在聊天框打的訊息 (如 prompt)

st.chat_message() 可以傳 name 參數給他，表示是誰講了這句話 – 只是為了 UI 顯示的目的，並沒有存起來 – 一些特殊的名字像 “user”, “assistant” 會在左邊顯示特別頭像

在 st.chat_message() 這個元件，可以用 Context Manager 的方式畫出訊息 – 也就是在他底下寫的程式都是在他的範圍情境內去畫。所以

    with st.chat_message("user"):
        st.markdown(prompt)

的意思就是「畫出一個 user 頭像的聊天訊息，而且訊息內容是 prompt 裡的值，以 markdown 格式畫出」

Rerun 重畫 #

當使用者在下面打字，聊天後你要把訊息畫在上面 – 也就是你想讓畫面更新的時候，Streamlit 的哲學是「把你的 python 程式重跑一遍」

因此，只靠上面的程式 st.chat_message() / st.write()，畫面上 UI 永遠只會顯示兩條訊息，沒有歷史

所以我們要做兩件事

把歷史訊息記錄起來
把所有訊息畫出來

# 只是個初始化
if "messages" not in st.session_state.keys():
    st.session_state.messages = []

# 重新畫所有訊息
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

if prompt := st.chat_input("任何魔方的問題都可以問我唷"):
    # 存訊息（給下次重畫用）
    st.session_state.messages.append({"role": "user", "content": prompt})
    ...
    st.session_state.messages.append({"role": "assistant", "content": response.response})

記錄可以用 st.session_state，就是個 key-value 的儲存，streamlit 每次重畫都還是會記得

Session State 每個瀏覽窗都會有自己的一個 session; 開兩窗不用怕打架

Streamlit 跟 LlamaIndex / RAG query engine 互動 #

在這個範例我沒有把 Llamaindex RAG 額外包成一個 web API，而是單純跟 Streamlit app 放在同一個 python 程式裡面

RAG 程式碼獨立放在 rag.py , 而 app.py 就只負責 UI

# Init LlamaIndex RAG query engine
if "query_engine" not in st.session_state.keys():
    # get_query_engine() 實作在 rag.py
    st.session_state.query_engine = get_query_engine()

query_engine = st.session_state.query_engine
with st.chat_message("assistant"):
    response = query_engine.query(prompt)

這邊要注意的是關於 LlamaIndex “index”：因為 UI 互動會重跑，每次都要載資料建 index 很累。除了儲存在各窗 session ，另一種想法是「存在 global 共用的 cache」

@st.cache_resource(show_spinner=False)
def load_index():
    if not os.path.exists(INDEX_PERSIST_DIR):
        raise IOError("Need to build the index first.")

    with st.spinner(text="Loading my knowledge..."):
        storage_context = StorageContext.from_defaults(persist_dir=INDEX_PERSIST_DIR)
        index = load_index_from_storage(storage_context)

    return index

利用 @st.cache_resource 把模型/index 存到 cache 裡，題外話 @st.cache_data 可存 serializable 物件/資料

這是 decorator，所以掛在回傳 index 的 function 上就好

（with st.spinner(...) 不用理他，那只是一開始在畫面上畫出「載入中」的圖案罷了）

執行 #

執行也很簡單，只要 streamlit run <python檔> 就好囉！

總之用 Streamlit 能夠很快速做出 POC 介面的部分，對於前端苦手的人來說很有幫助！

若您覺得有趣, 請追蹤我的Facebook 或 Linkedin, 讓你獲得更多資訊！