LangChain 怎麼玩？如何做資料庫查詢機器人

Posted on Mar 14, 2024 in LangChain , Python 程式設計 - 高階 by Amo Chen ‐ 5 min read

覺得我們的內容實用嗎？ MyApollo 電子報讀者募集中！歡迎訂閱電子報!

作為軟體開發者，你開始發想任何語言模型可以應用的領域了嗎？

如果你沒有任何想法，也許可以先看看怎麼藉由語言模型，讓我們可以用對話方式，就能夠自動對資料庫執行 SQL 進行查詢，而且過程中完全不需要寫任何 SQL 喔！

也許看完本文，你會對語言模型的應用想像會有更清楚的畫面！

本文環境

$ pip install langchain SQLAlchemy

本文需要 Ollama 指令與 Meta 公司提供的 llama2 模型(model)，ollama 指令請至 Ollama 官方頁面下載安裝，安裝完成之後即可執行指令安裝 llama2 模型，其安裝指令如下：

$ ollama run llama2

p.s. 執行上述模型需要至少 8GB 記憶體， 3.8G 硬碟空間

除了 llama2 模型之外，本文還需要 SQLite 3 版本的 Northwind 範例資料庫。

Northwind 範例資料庫是 Microsoft Access 2000 所提供的練習用的資料庫，不過本文作為教學文章，為求環境容易搭建，故使用 SQLite3 版本（下載連結），下載之後請與範例程式放在同一資料夾內，例如：

.
├── northwind.db
└── query_bot.py

如對 LangChain 框架不熟者，建議先看過本系列教學文後再閱讀本篇。

資料庫查詢聊天機器人

通常商業數據分析人員在分析資料時，多數會從資料庫進調出感興趣的資料後，再進一步進行分析，所以看到數據分析師職缺要求上出現「懂 SQL 」這一項並不稀奇，這多半代表此職缺擁有操作資料庫的權限，不過，並不是每個數據分析師都知道如何使用 SQL 。

受益於語言模型對於自然語言的理解/生成能力，學會使用 SQL 對於數據分析師而言，可能不再是一項必備技能，因為藉由自然語言與語言模型進行對話，進而產生 SQL 並實際於資料庫中執行，最後生成答案，在如今是完全可以做到的一件事，甚至未來可能會更加簡單、直覺，譬如可以想像得到未來資料庫的 Client APP 會內建語言模型。

總之，一起來看看如何能做這件事吧！

流程圖

以下是資料庫查詢聊天機器人的運作流程（為求容易理解，該圖經過簡化）：

簡單說明一下，上圖先後使用 2 次 Prompt：

第 1 次，藉由輸入問題，產生欲執行的 SQL, 也就是圖中藍色虛線框的部分，不過語言模型在不知道我們資料庫有哪些表格以及表格有哪些欄位的情況下，是無法生成相對應的 SQL 的，所以我們仍然需要在 prompt 中附上 DB schema 作為 context, 給語言模型作為參考，以生成 SQL 。
第 2 次，在得到使用者問題以及用第 1 次 prompt 生成的 SQL 之後，仍需要先將 SQL 送到資料庫執行，藉此得到 SQL 執行後的結果，最終將使用者的問題、 SQL 以及 SQL 執行結果作為第 2 次 prompt 所需要的參數，讓語言模型生成回應，以作為最終使用者看到的回答。

完整程式碼

以下是資料庫查詢聊天機器人的完整程式碼，重點部分容後解說：

from langchain_community.llms import Ollama
from langchain_community.utilities import SQLDatabase
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from sqlalchemy.exc import OperationalError

llm = Ollama(model='llama2')


db = SQLDatabase.from_uri("sqlite:///./northwind.db")


def get_db_schema(_):
    return db.get_table_info()


def run_query(query):
    try:
        return db.run(query)
    except (OperationalError, Exception) as e:
        return str(e)


gen_sql_prompt = ChatPromptTemplate.from_messages([
    ('system', 'Based on the table schema below, write a SQL query that would answer the user\'s question: {db_schema}'),
    ('user', 'Please generate a SQL query for the following question: "{input}". \
     The query should be formatted as follows without any additional explanation: \
     SQL> <sql_query>\
    '),
])


class SqlQueryParser(StrOutputParser):
    def parse(self, s):
        r = s.split('SQL> ')
        if len(r) > 0:
            return r[1]
        return s


gen_query_chain = (
    RunnablePassthrough.assign(db_schema=get_db_schema)
    | gen_sql_prompt
    | llm
    | SqlQueryParser()
)


gen_answer_prompt = ChatPromptTemplate.from_template("""
Based on the provided question, SQL query, and query result, write a natural language response.
No additional explanations should be included.

Question: {input}
SQL Query: {query}
Query Result: {result}

The response should be formatted as follows:
'''
Executed: {query}
Answer: <answer>
'''
""")


chain = (
    RunnablePassthrough.assign(query=gen_query_chain).assign(
        result=lambda x: run_query(x["query"]),
    )
    | gen_answer_prompt
    | llm
)

input_text = input('>>> ')
while input_text.lower() != 'bye':
    if input_text:
        response = chain.invoke({
            'input': input_text,
        })
        print(response)
    input_text = input('>>> ')

上述程式執行之後，可以問它有幾位員工看看：

>>> How many employees are there?
Executed: SELECT COUNT(*) FROM Employees;
Answer: There are 9 employees.

它會告訴我們有 9 位員工，實際以 SQL SELECT COUNT(*) FROM Employees; 執行，也可以得到數字為 9 。

除此之外，也可以問問其他問題，例如問它產品的最高單價：

>>> What's the highest unit price of a product?
Executed: SELECT MAX(UnitPrice) FROM Products;
Answer: 263.5

實際執行 SQL SELECT MAX(UnitPrice) FROM Products; 也可以得到數字為 263.5 。

這就是 1 個資料庫查詢機器人的基本範例。

關鍵部分解說

接下來，解說上述程式碼的關鍵部分。

首先，先載入資料庫以及定義需要的 2 個函數，一為取得資料庫表格結構，一為執行 SQL 之用，此處可以利用 LangChain 提供的 SQLDatabase 類別，它的底層其實是 SQLAlchemy:

db = SQLDatabase.from_uri("sqlite:///./northwind.db")


def get_db_schema(_):
    return db.get_table_info()


def run_query(query):
    try:
        return db.run(query)
    except (OperationalError, Exception) as e:
        return str(e)

接著，我們需要定義 1 個 prompt 告訴語言模型產生 SQL, 如流程圖說明所述，我們也需要：

附上資料庫有哪些表格、表格有哪些欄位的資訊，讓語言模型知道有哪些表格、欄位可以使用，對應以下 prompt 的 db_schema
使用者想查詢的問題，對應以下 prompt 的 input

此外，我們希望語言模型將 SQL 按照 SQL> <答案> 的格式輸出，而且不需要為我們講解 SQL 的說明，對應 The query should be formatted as follows without any additional explanation ... 的部分：

gen_sql_prompt = ChatPromptTemplate.from_messages([
    ('system', 'Based on the table schema below, write a SQL query that would answer the user\'s question: {db_schema}'),
    ('user', 'Please generate a SQL query for the following question: "{input}". \
     The query should be formatted as follows without any additional explanation: \
     SQL> <sql_query>\
    '),
])

有生成 SQL 的 prompt 之後，再定義 1 個 Output Parser, 幫助我們從該 prompt 取出 SQL, 由於我們已經規定固定生成格式，所以可以簡單地用 split() 函式取得 SQL:

class SqlQueryParser(StrOutputParser):
    def parse(self, s):
        r = s.split('SQL> ')
        if len(r) > 0:
            return r[1]
        return s

重要的是，在使用者輸入 input 之後，變成 prompt 之前，我們必須提供 db_schema 才能正確組成 prompt, 因此在 gen_sql_prompt 之前，我們多加 1 個 RunnablePassthrough.assign(db_schema=get_db_schema) ，指定要加 1 個 db_schema 的值，該值需要呼叫 get_db_schema() 函式，如此就可以把資料庫結構一起放進 prompt 之中：

gen_query_chain = (
    RunnablePassthrough.assign(db_schema=get_db_schema)
    | gen_sql_prompt
    | llm
    | SqlQueryParser()
)

到這一步，就完成流程圖中藍色虛線框的部分。

再來，定義第 2 個 prompt, 讓語言模型根據使用者輸入、 SQL 以及 SQL 執行結果產生答案：

gen_answer_prompt = ChatPromptTemplate.from_template("""
Based on the provided question, SQL query, and query result, write a natural language response.
No additional explanations should be included.

Question: {input}
SQL Query: {query}
Query Result: {result}

The response should be formatted as follows:
'''
Executed: {query}
Answer: <answer>
'''
""")

最後把剩下的 chain 組合起來，也就是藍色虛線框以及其他部分：

chain = (
    RunnablePassthrough.assign(query=gen_query_chain).assign(
        result=lambda x: run_query(x["query"]),
    )
    | gen_answer_prompt
    | llm
)

上述 chain 比較令人費解的部分如下：

RunnablePassthrough.assign(query=gen_query_chain).assign(
    result=lambda x: run_query(x["query"]),
)

說明一下，當我們呼叫 chain.invoke({'input': 'question'}) 時，上述的 RunnablePassthrough.assign(query=gen_query_chain) 會先執行， {'input': 'question'} 會被丟到為 gen_query_chain 進而產生 SQL, 這個產生的 SQL 會被指定為 key query 的值，最後變成：

{
    'input': 'question',
    'query': '....SQL...',
}

接著，上述結果又被傳給下一個呼叫 .assign(result=lambda x: run_query(x["query"]))，所以 lambda x: run_query(x["query"]) 中的 x 才有 key query 可用，最後 run_query(x["query"]) 的執行結果又被指定為 key result 的值：

{
    'input': 'question',
    'query': '....SQL...',
    'result': '...',
}

畫成圖的話，會更清楚，實際上是執行 2 次 RunnablePassthrough 的效果：

如此就準備好組成第 2 個 prompt 以產生最終答案所需要的 3 個參數了！

剩下的部分就都是 chain 的部分了，也就是 prompt -> LLM -> answer 的步驟。

總結

本文展示如何使用語言模型結合資料庫查詢的功能，讓我們能夠透過文字聊天對話的方式就可以取得資料庫中的數據。

不過，本文的範例也僅僅只是揭開此類型應用的核心運作原理而已，仔細玩過幾次的話，就會發現並不是每次 SQL query 都可以使用最有效率的方式達成，這些問題都是屬於在做這類型應用要額外處理的問題，甚至也必須處理資料庫的安全問題，例如給予最小訪問權限、限定讀取(read only)、限制可存取的表格、欄位等等，要打造一款安全、易用的語言模型應用，還是需要了解、注意各種面向的知識的。

以上！

Enjoy!

References

SQLite3 version of Microsoft’s Northwind Database

LangChain - Querying a SQL DB

覺得我們的內容實用嗎？ MyApollo 電子報讀者募集中！歡迎訂閱電子報!

python langchain llm

LangChain 怎麼玩？如何做資料庫查詢機器人

本文環境

資料庫查詢聊天機器人

流程圖

完整程式碼

關鍵部分解說

總結

References

對抗久坐職業傷害

贊助我們的創作

LangChain 怎麼玩？如何做資料庫查詢機器人

本文環境 #

資料庫查詢聊天機器人 #

流程圖 #

完整程式碼 #

關鍵部分解說 #

總結 #

References #

對抗久坐職業傷害

贊助我們的創作

你可能也會感興趣的文章

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

LangChain 怎麼玩？用 Hugging Face 結合 LangChain，做個英文對話機器人模擬面試吧

LangChain 怎麼玩？用 LangSmith 幫忙追查問題

本文環境

資料庫查詢聊天機器人

流程圖

完整程式碼

關鍵部分解說

總結

References