LangChain 怎麼玩？用 Streamlit 打造 LLM 個人工具箱
Last updated on Sep 24, 2024 in LangChain , Python 程式設計 - 高階 by Amo Chen ‐ 6 min read
每個程式設計師多多少少都有打造私人工具箱，不僅可以增加工作效率，還可以跟著職業生涯累積起來帶著走。
AI 時代來臨，工具箱當然免不了要多一些 AI 相關的工具，本文將介紹怎麼用 LangChain 結合 Python 知名套件 Streamlit 打造屬於你的個人工具箱！
本文環境
$ pip install streamlit langchain
本文需要 Ollama 指令與 Meta 公司提供的 llama2 模型(model)，ollama 指令請至 Ollama 官方頁面下載安裝，安裝完成之後即可執行指令安裝 llama2 模型，其安裝指令如下：
$ ollama run llama2
p.s. 執行上述模型需要至少 8GB 記憶體， 3.8G 硬碟空間
Streamlit 簡介
Streamlit 是 1 個 Python 知名套件，主要用於建立資料應用程式(data apps)，因此很多與資料分析、視覺化相關的工作，都可以使用 Streamlit 實作，整體開發過程就像在寫 Python script 一樣輕鬆。
Streamlit 訴求簡單、易用，因此 Streamlit 內建許多 UI 元件，如文字框、按鈕、圖表和地圖等等，並且整合 pandas, Matplotlib, Plotly 等 Python 常用套件，就算我們沒有任何前端技能，我們也能夠用少少的程式碼就做出不錯的 data app 。
而隨著 LLM 相關應用的盛行，因此 Streamlit 也提供對話(chat)相關的 UI 元件，所以使用 Streamlit 開發 LLM 相關的個人專屬應用也很合適。
簡單的 Streamlit + LangChain 應用 — Ask Me Anything
進入主題，先做個簡單的 Chat 應用跟語言模型對話。
1 個最簡單的對話應用，就只需要 1 個表單、 1 個文字輸入框、 1 個送出按鈕、 1 個顯示回應的 UI 元件，所以我們可以從 Streamlit 的 API reference 挑出 4 個合適的 UI 元件：
- st.text_area 文字輸入框
- st.form_submit_button 送出按鈕
- st.form 表單的容器，與送出按鈕是成對的，如果要使用送出按鈕，就必須在 st.form 內使用
- st.info 用以顯示回應的文字，也可以改用 st.text 等 UI 元件
如同前文所述，打造 Streamlit 的應用就像在寫 Python script 一樣，通常需要什麼元件就直接寫在 script 裡即可，所以集合上述 4 種 UI 元件的框架程式碼如下：
import streamlit as st
with st.form('form'):
text = st.text_area('Enter text:', '')
submitted = st.form_submit_button('Submit')
if submitted:
st.info('Hello World')
上述程式需要以 streamlit run <streamlit-script>.py 指令執行，例如：
$ streamlit run example.py
執行成功之後，瀏覽器會自動打開，我們將可以看到如下畫面：
Streamlit 是否很簡單呢？至於其他 UI 元件如何使用請參考官方文件。
接著，就只要在前述框架上，將 LangChain 與語言模型放進去即可，其程式碼如下：
import streamlit as st
from langchain_community.llms import Ollama
llm = Ollama(model='llama2')
def generate_response(text):
return llm.invoke(text):
st.title('Ask Me Anything')
with st.form('form'):
text = st.text_area('Enter text:', '')
submitted = st.form_submit_button('Submit')
if submitted:
st.info(generate_response(text))
上述程式先指定使用 llama2 語言模型，接著定義 1 個名稱為 generate_response 的函式，接受傳入 prompt 並取得生成的回應，使用者輸入的 prompt 則是從 st.text_area('Enter text:', '') 取得，並在使用者點擊 Submit 按鈕之後，呼叫 generate_response() 並將生成的回應顯示於 st.info() 元件之中。
上述程式碼同樣需要以 streamlit run <streamlit-script>.py 指令執行，例如：
$ streamlit run <streamlit-script>.py
其執行畫面如下圖所示：
語言模型回應好慢好煩躁，怎麼辦？
使用前述範例一定會遇到 1 個問題，那就是每次輸入完都要等一段時間，讓語言模型生成結束才會顯示回應，使用體驗肯定相當煩躁。
為此， LangChain 的 Runnable 其實早有提供 stream() 函式，可以不用等生成結束，只要一有部分回應就可以先將結果送出，體驗上會與 ChatGPT 相同，一直有文字輸出直到結束，體驗上相對好很多。
使用 stream() 的方式很簡單，只要使用 for 迴圈一直 iterate 它即可：
for r in llm.stream(input):
print(r)
如果完整的回應是 Hello, nice to see you, 那 stream() 函式就會拆成 token 逐個輸出，例如 Hello > nice to > see > you 的順序（p.s. token 拆分依語言模型不同有不同結果）。
改成使用 stream() 的 Streamlit 程式碼如下：
import streamlit as st
from langchain_community.llms import Ollama
llm = Ollama(model='llama2')
def generate_response(text):
for r in llm.stream(text):
yield r
st.title('Ask Me Anything')
with st.form('form'):
text = st.text_area('Enter text:', '')
submitted = st.form_submit_button('Submit')
if submitted:
resp = ""
for r in generate_response(text):
resp += r
st.info(resp)
如果執行上述程式碼，並試圖對話的話，將會發現其結果如下所示，文字並不像 ChatGPT 那樣逐個顯示：
這是因為 st.info() 並不支援 streming 形式的輸出。
為了解決上述問題，我們需要改用 Streamlit 為了 Chat 應用所開發的 st.write_stream() UI 元件，它支援接受 generator, iterable 或是 stream-like sequence, 能夠做到像 ChatGPT 一樣逐個顯示的功能。
改使用 st.write_stream() UI 元件的範例如下：
import streamlit as st
from langchain_community.llms import Ollama
llm = Ollama(model='llama2')
def generate_response(text):
for r in llm.stream(text):
yield r
st.title('Ask Me Anything')
with st.form('my_form'):
text = st.text_area('Enter text:', '')
submitted = st.form_submit_button('Submit')
if submitted:
st.write_stream(generate_response(text))
上述程式執行之後，並再次對話之後，將可以看到文字逐個顯示：
更多關於 Chat 應用的元件請參考 Stream - Chat elements。
打造個人工具箱(toolbox)
理解如何將 LangChain 與 Streamlit 結合之後，我們就具備打造個人工具箱的能力了！
以下是本文打造的工具箱介面：
上述工具箱相當簡單，內建 2 個 AI 角色，以及幾種常用的 prompt 樣版(template)，用以減少重複輸入 prompt 的情況，額外用到 2 種 UI 元件：
- st.selectbox 選擇角色
- st.radio 選擇樣版
與 LangChain 整合的部分，則是跟前文所差無幾。
其完整程式碼如下：
import streamlit as st
from langchain_community.llms import Ollama
from langchain_core.prompts import ChatPromptTemplate
ROLE_PYTHON_DEVELOPER = 'Senior Python Developer'
TMPL_NAMING = 'Naming'
TMPL_CODE_REVIEW = 'Code Review'
ROLE_ENGLISH_TEACHER = 'Professional English Teacher'
TMPL_CORRECT_GRAMMAR = 'Correct English Grammar'
TMPL_CORRECT_GRAMMAR_WITH_DESCR = 'Correct English Grammar (with additional description)'
TMPL_NONE = 'None'
llm = Ollama(model='llama2')
st.title('Ask Me Anything')
role = st.selectbox(
"Which AI role would you like to ask?",
(ROLE_PYTHON_DEVELOPER, ROLE_ENGLISH_TEACHER),
index=0,
)
prompt = ChatPromptTemplate.from_messages([
(
"system",
"You are a {role}, you job is answering user's question.",
),
("user", "{input}"),
])
chain = prompt | llm
def generate_response(text):
for r in chain.stream({'input': text, 'role': role}):
yield r
template = TMPL_NONE
if role == ROLE_PYTHON_DEVELOPER:
template = st.radio(
"Templates",
[
TMPL_NAMING,
TMPL_CODE_REVIEW,
TMPL_NONE,
],
horizontal=True
)
elif role == ROLE_ENGLISH_TEACHER:
template = st.radio(
"Templates",
[
TMPL_CORRECT_GRAMMAR_WITH_DESCR,
TMPL_CORRECT_GRAMMAR,
TMPL_NONE,
],
horizontal=True
)
def use_template(text):
if template == TMPL_NONE:
return text
if template == TMPL_NAMING:
return f'Is "{text}" a good variable name in Python?'
if template == TMPL_CODE_REVIEW:
return (
"Please assist me in reviewing the following code snippet. "
"Two hard rules apply: "
"1. Function names and variables must be clear. "
"2. There should be no performance issues. "
"If the code appears clear and efficient, simply respond with 'LGTM'. "
"Otherwise, please identify any issues and provide additional explanation. "
"The code snippet is:\n"
f"'''\n{text}\n'''"
)
if template == TMPL_CORRECT_GRAMMAR_WITH_DESCR:
return (
"Please assist me in correcting the grammar issues. "
"The text is:\n"
f"'''\n{text}\n'''"
)
if template == TMPL_CORRECT_GRAMMAR:
return (
"Please assist me in correcting the grammar issues. "
"You don't need to provide additional explanation. "
"The text is:\n"
f"'''\n{text}\n'''"
)
with st.form('form'):
text = st.text_area('Enter text:', '')
submitted = st.form_submit_button('Submit')
if submitted:
st.write_stream(generate_response(use_template(text)))
上述程式碼主要多了 use_template() 函式，其作用在於判別使用者選了哪個樣版，並將使用者的輸入放到樣版之中，再交給語言模型產生回應。
其執行結果如下所示：
用 Streamlit 打造個人 AI 工具箱是不是簡單、直覺很多呢？
快取資源
Streamlit runs your script from top to bottom at every user interaction or code change.
目前的程式碼存在 1 個問題，那就是每次重新渲染(re-render) UI 時，都會造成重新執行 llm = Ollama(model='llama2')，重新載入語言模型會影響重新渲染的速度。
所以我們需要使用快取減少重新載入語言模型的情況，因此可以使用 st.cache_resource 將 Ollama(model='llama2') 快取起來，修改過後的程式碼如下：
import streamlit as st
from langchain_community.llms import Ollama
from langchain_core.prompts import ChatPromptTemplate
ROLE_PYTHON_DEVELOPER = 'Senior Python Developer'
TMPL_NAMING = 'Naming'
TMPL_CODE_REVIEW = 'Code Review'
ROLE_ENGLISH_TEACHER = 'Professional English Teacher'
TMPL_CORRECT_GRAMMAR = 'Correct English Grammar'
TMPL_CORRECT_GRAMMAR_WITH_DESCR = 'Correct English Grammar (with additional description)'
TMPL_NONE = 'None'
@st.cache_resource
def load_llm():
return Ollama(model='llama2')
llm = load_llm()
st.title('Ask Me Anything')
role = st.selectbox(
"Which AI role would you like to ask?",
(ROLE_PYTHON_DEVELOPER, ROLE_ENGLISH_TEACHER),
index=0,
)
prompt = ChatPromptTemplate.from_messages([
(
"system",
"You are a {role}, you job is answering user's question.",
),
("user", "{input}"),
])
chain = prompt | llm
def generate_response(text):
for r in chain.stream({'input': text, 'role': role}):
yield r
template = TMPL_NONE
if role == ROLE_PYTHON_DEVELOPER:
template = st.radio(
"Templates",
[
TMPL_NAMING,
TMPL_CODE_REVIEW,
TMPL_NONE,
],
horizontal=True
)
elif role == ROLE_ENGLISH_TEACHER:
template = st.radio(
"Templates",
[
TMPL_CORRECT_GRAMMAR_WITH_DESCR,
TMPL_CORRECT_GRAMMAR,
TMPL_NONE,
],
horizontal=True
)
def use_template(text):
if template == TMPL_NONE:
return text
if template == TMPL_NAMING:
return f'Is "{text}" a good variable name in Python?'
if template == TMPL_CODE_REVIEW:
return (
"Please assist me in reviewing the following code snippet. "
"Two hard rules apply: "
"1. Function names and variables must be clear. "
"2. There should be no performance issues. "
"If the code appears clear and efficient, simply respond with 'LGTM'. "
"Otherwise, please identify any issues and provide additional explanation. "
"The code snippet is:\n"
f"'''\n{text}\n'''"
)
if template == TMPL_CORRECT_GRAMMAR_WITH_DESCR:
return (
"Please assist me in correcting the grammar issues. "
"The text is:\n"
f"'''\n{text}\n'''"
)
if template == TMPL_CORRECT_GRAMMAR:
return (
"Please assist me in correcting the grammar issues. "
"You don't need to provide additional explanation. "
"The text is:\n"
f"'''\n{text}\n'''"
)
with st.form('form'):
text = st.text_area('Enter text:', '')
submitted = st.form_submit_button('Submit')
if submitted:
st.write_stream(generate_response(use_template(text)))
總結
Streamlit 憑藉其直覺、易用的應用開發方式，在 Python 的生態系中佔有一席之地，不僅能用來開發資料應用，其所提供的 Chat 相關 UI 元件也相當適合開發 LLM 應用。
本文作為 LangChain 教學文，僅專注於如何將 Streamlit 與 LangChain 進行整合，但實際上 Streamlit 也有相當多的元件能夠解決各種分析需求，有興趣的人可以至 Streamlit 官網一探究竟。
以上！
Enjoy!
References
Build an LLM app using LangChain
Playing with Streamlit and LLMs.
