LangChain 怎麼玩？用 Streamlit 打造 LLM 個人工具箱

每個程式設計師多多少少都有打造私人工具箱，不僅可以增加工作效率，還可以跟著職業生涯累積起來帶著走。

AI 時代來臨，工具箱當然免不了要多一些 AI 相關的工具，本文將介紹怎麼用 LangChain 結合 Python 知名套件 Streamlit 打造屬於你的個人工具箱！

本文環境 #

• macOS
• Python 3
• Streamlit
• LangChain
• Ollama

$ 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 元件：

1. st.text_area 文字輸入框
2. st.form_submit_button 送出按鈕
3. st.form 表單的容器，與送出按鈕是成對的，如果要使用送出按鈕，就必須在 st.form 內使用
4. 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 元件：

1. st.selectbox 選擇角色
2. 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 憑藉其直覺、易用的應用開發方式，在 Python 的生態系中佔有一席之地，不僅能用來開發資料應用，其所提供的 Chat 相關 UI 元件也相當適合開發 LLM 應用。

本文作為 LangChain 教學文，僅專注於如何將 Streamlit 與 LangChain 進行整合，但實際上 Streamlit 也有相當多的元件能夠解決各種分析需求，有興趣的人可以至 Streamlit 官網一探究竟。

以上！

Enjoy!

References #

Streamlit - API Reference

Build an LLM app using LangChain

Playing with Streamlit and LLMs.