使用 TiDB Vector 搭建 RAG 應用 - TiDB 文件問答小助手

balahoho發表於2024-06-04
TiDB

本文首發至TiDB社群專欄：https://tidb.net/blog/7a8862d5

前言

繼上一次《TiDB Vector搶先體驗之用TiDB實現以圖搜圖》後，就迫不及待的想做一些更復雜的應用。上一篇在 TiDB 社群專欄釋出以後還是有很多社群朋友不明白向量的應用場景到底是什麼，這次用一個更直觀的場景來體現向量檢索在 AI 應用開發的重要性。

知識庫問答是目前 AGI 領域應用最多的場景之一，本次我基於 TiDB Vector 給 TiDB 搭建一個文件問答小助手。

前置知識

上一篇提到把非結構化資料轉化為向量表示需要用到 embedding 模型，但這種模型和大家所瞭解的 GPT 大語言模型（LLM）還不太一樣，他們有著不同的作用。

text-embedding-ada-002 和 GPT（Generative Pre-trained Transformer）是兩種不同型別的模型，它們在設計和功能上有所不同，但都由 OpenAI 推出。

  1. text-embedding-ada-002：這是一種文字嵌入模型，它的主要功能是將文字轉換為高維向量表示（嵌入）。這種嵌入可以捕捉文字的語義和語境資訊，通常用於文字相似度計算、推薦系統等任務中。text-embedding-ada-002 使用了 AdaIN（Adaptive Instance Normalization）技術，透過學習將文字對映到高維向量空間中。
  2. GPT（Generative Pre-trained Transformer）：GPT 是一系列基於 Transformer 架構的語言模型，由 OpenAI 釋出。這些模型被設計用於自然語言處理任務，如文字生成、文字理解、問答等。GPT 模型在預訓練階段使用了大規模的文字資料，然後可以在各種任務上進行微調或直接應用。

雖然 text-embedding-ada-002 和 GPT 都與文字處理相關，但它們的功能和用途不同。text-embedding-ada-002 主要用於文字嵌入，而 GPT 則是一個通用的自然語言處理模型，可以用於多種文字相關任務。因此，它們之間的關係是它們都由 OpenAI 推出，並且都是用於文字處理的模型，但它們的具體功能和設計是不同的。

注：以上解釋由 ChatGPT 生成 。

本次實驗中我用text-embedding-ada-002對 TiDB 的中文文件做 embedding 轉化存入到 TiDB Serverless，用 GPT-4 生成最終的問題答案。

為什麼需要 RAG

在各種各樣的資訊渠道，相信大家已經被 RAG 這個詞在視覺上轟炸了很長時間，但是我估計大部分 DBA 看了依然不明白到底是什麼，我爭取用一個小例子來講明白。

眾所周知，大語言模型都是預訓練模型，也就是說他們的語料是不會及時更新的，類似於資料庫裡的snapshot。比如我去問 ChatGPT 目前（24年5月14號） TiDB 的最新版本是多少，它的回答一定不讓人滿意。

企業微信截圖_20240508152948.png

如果我就想讓 ChatGPT 告訴我 TiDB 的最新版本號，通常有兩種辦法：

  • 全量模式：OpenAI 把 GPT 重新訓練一次，生成新的 snapshot（超級燒錢）
  • 增量模式：給 GPT 投餵一些新的資訊，比如在問題的上下文中把新版本資訊告訴 GPT，讓它根據提供的上下文來回答

很明顯第二種方法更切實際，工程實現上更容易，成本也更低。

我們從指定的文件、甚至是搜尋引擎中找到相關的資訊丟給 GPT ，藉助 GPT 的推理能力並且限定它的上下文內容得到最終答案。透過臨時餵給 LLM 的資料提升它的能力，這就構成了 RAG 的靈魂：檢索（Retrieval）、增強（Augmented）、生成（Generation）。

這就是一個最基礎的 RAG （也稱之為樸素 RAG，Native RAG）流程，在此基礎上如果繼續最佳化提升準確度的話還可以引入 Rerank 等相關技術形成高階 RAG（Advance RAG）。

企業微信截圖_20240508155447.png

可以發現檢索是 RAG 裡非常重要的一個流程，因此 TiDB 的向量檢索能力就能起到關鍵作用。目前市面上見到的絕大部分 AI 應用，都是用 RAG 架構來搭建的。

到這裡不知道大家會不會有個疑問：

既然檢索（Retrieval）就能得到想要的答案，為什麼要多此一舉再問一遍 LLM ？

TiDB 知識庫問答小助手

基於前面介紹的 RAG 架構，下面我逐漸用程式碼實現讓 GPT 能回答剛才那個 TiDB 版本號問題。

1、文件切分和向量化

LangChain 官方已經對 TiDB Vector 做了整合，藉助 LangChain 的 vectorstore 元件能夠對 TiDB Vector 實現高效操作。

from langchain_community.vectorstores import TiDBVectorStore

剛好我本地有一份之前下載的 TiDB v7.6.0 PDF 文件，先把檔案處理後儲存到 TiDB Serverless 中：

import os
from langchain_community.document_loaders import PyPDFLoader
from langchain_openai import AzureOpenAIEmbeddings

os.environ["OPENAI_API_TYPE"] = "azure"
os.environ["OPENAI_API_VERSION"] = "2024-02-01"
os.environ["AZURE_OPENAI_ENDPOINT"] = "https://heao-ai-test1.openai.azure.com/"
os.environ["OPENAI_API_KEY"] = "xxxxx"

embeddings = AzureOpenAIEmbeddings(model="text-embedding-ada-002")

TIDB_CONN_STR="mysql+pymysql://xxxx.root:xxxx@gateway01.eu-central-1.prod.aws.tidbcloud.com:4000/test?ssl_ca=C:\\Users\\59131\\Downloads\\isrgrootx1.pem&ssl_verify_cert=true&ssl_verify_identity=true"
TABLE_NAME = "semantic_embeddings"

def load_docment(path):
     loader = PyPDFLoader(path)
     chunks = loader.load_and_split()
     db = TiDBVectorStore.from_documents(
         documents = chunks,
         embedding = embeddings,
         table_name = TABLE_NAME,
         connection_string = TIDB_CONN_STR,
         distance_strategy = "cosine",  # default, another option is "l2"
     )

簡單幾行程式碼 LangChain 幫我們把 split、embedding、入庫全部都做了，再看一下 TiDB Vector 中生成的表結構：

CREATE TABLE `semantic_embeddings` (
`id` varchar(36) NOT NULL,
`embedding` vector<float>(1536) NOT NULL COMMENT 'hnsw(distance=cosine)',
`document` text DEFAULT NULL,
`meta` json DEFAULT NULL,
`create_time` datetime DEFAULT CURRENT_TIMESTAMP,
`update_time` datetime DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
PRIMARY KEY (`id`) /*T![clustered_index] CLUSTERED */
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin

這裡做 chunk 用最簡單粗暴的形式，每頁文件為一個 chunk，chunk 的文字內容存入 document欄位，向量化後的內容存入embedding欄位，這是一個1536維的向量同時建立了hnsw索引，meta欄位儲存的是 chunk 的一些元資訊，如檔名、頁號等。

企業微信截圖_20240514155507.png

文件預處理和 chunk 劃分方式對召回質量有很大的影響，這些屬於 RAG 調優範疇，簡單起見本文不做討論，重點關注整體實現流程。

2、向量檢索召回

知識庫準備好以後就可以根據我們提出的問題在語義層面搜尋相關內容，主要依賴 TiDB 的向量檢索能力，這一步稱為召回。

def get_tidb_connenction():
    db = TiDBVectorStore.from_existing_vector_table(
        embedding = embeddings,
        table_name = TABLE_NAME,
        connection_string = TIDB_CONN_STR,
        distance_strategy = "cosine", 
    )
    return db

def retrieval_from_tidb(db, query):
    docs_with_score = db.similarity_search_with_score(query, k=3)
    context = ""
    for doc, score in docs_with_score:
        context += doc.page_content + "\n"
    return context

我們去 TiDB 中查詢到相似度最高的 TOP 3資訊，簡單拼接後組裝成上下文返回。

3、構建 Prompt

Prompt 是和 LLM 溝通的語言，透過 Prompt 我們可以引導大模型控制它的輸出和準確度。就好比我們要去搜素引擎或 github上搜想要的內容時，如何提問是一門藝術。它有一套非常全面的方法論，這裡不做過多贅述，用常用的格式構建一個 Prompt 即可：

def build_prompt(context, question):
    template = """你的角色是一個TiDB知識庫文件助手，希望你能幫我解決使用TiDB遇到的問題。我會給你一些輔助上下文資訊來幫助你回答問題，如果你不知道答案，就告訴我抱歉無法回答，不要回答其他不確定的內容。
    <context>
    {context}
    </context>
    <question>
    {question}
    </question>
    """
    prompt = ChatPromptTemplate.from_template(template)
    value = prompt.format_prompt(question=question, context=context)

    return value

4、LLM 結果生成

下一步將得到的 Prompt 傳給大模型並獲得返回結果，再把整個 RAG 的呼叫鏈串起來封裝成rag_invoke方法：

def get_answer(prompt):
    llm = AzureChatOpenAI(
        temperature=0.0,
        deployment_name="gpt-4", #上面的deployment name
        model_name="gpt-4" #deployment 對應的model
    )
    response = llm.invoke(prompt)
    return response.content

def rag_invoke(question):
    db = get_tidb_connenction()
    context = retrieval_from_tidb(db,question)
    prompt = build_prompt(context,q)
    a = get_answer(prompt)
    print(f"RAG：{a}")

為了方便對比，我把不使用 RAG 直接呼叫 GPT API 的結果也列印出來：

def llm_invoke(question):
    llm = AzureChatOpenAI(
        temperature=0.0,
        deployment_name="gpt-4",
        model_name="gpt-4" 
    )
    print(f"GPT：{llm.invoke(question).content}")

5、效果演示

最後來看一下效果怎麼樣：

if __name__ == '__main__':
    # load_docment("C:/Users/59131/Downloads/tidb-dev-zh-manual.pdf")

    q = "TiDB的最新版本是多少？"
    print(f"Q: {q} \n")
    llm_invoke(q)
    print("-------------------------------")
    rag_invoke(q) 

E:\GitLocal\AITester>python langchain_doc.py
Q: TiDB的最新版本是多少？ 

GPT：截至我回答這個問題的時間（2021年11月），TiDB的最新穩定版本是5.2.1，釋出於2021年10月22日。但請注意，軟體的版本更新非常快，建議去官方網站檢視最新版本。
-------------------------------
RAG：根據提供的資訊，TiDB的最新版本是7.6.0-DMR，釋出日期為2024-01-25。

經過“增強”後，GPT 能準確的答出最新版本是7.6.0-DMR，看起來變得更聰明瞭，“增強”兩個字型現的恰到好處。

比如我再問一些 TiDB 比較新的特性：

E:\GitLocal\AITester>python langchain_doc.py
Q: TiDB中資源管控的單位是什麼？ 

GPT：TiDB中資源管控的單位是SQL查詢。
-------------------------------
RAG：TiDB中資源管控的單位是Request Unit (RU)。

標準版 GPT 甚至開始胡言亂語了。

前面提到為什麼生成答案還要再呼叫一次 LLM ，不直接使用 TiDB Vector 中返回的結果？以最初的 TiDB 最新版本問題為例，我們看一下向量檢索的結果是什麼，列印context變數的值：

TiDB
Binlog
版本TiDB
版本 說明
Local TiDB
1.0及
更低
版本
Kafka TiDB
1.0 ~
TiDB
2.1
RC5TiDB
1.0支
持
local
版本
和
Kafka
版本
的
TiDB
Bin-
log。
Cluster TiDB
v2.0.8-
binlog ，
TiDB
2.1
RC5
及更
高版
本TiDB
v2.0.8-
binlog
是一
個支
持
Clus-
ter版
本
TiDB
Binlog
的2.0
特殊
版本。
13.10.6.2.1 升級流程
注意：
如果能接受重新導全量資料，則可以直接廢棄老版本，按 TiDB Binlog 叢集部署 中的步驟重新部
署。
2279
16.2 TiDB版本釋出時間線
本文列出了所有已釋出的 TiDB版本，按釋出時間倒序呈現。
版本 釋出日期
6.5.8 2024-02-02
7.6.0-DMR 2024-01-25
6.5.7 2024-01-08
7.1.3 2023-12-21
6.5.6 2023-12-07
7.5.0 2023-12-01
7.1.2 2023-10-25
7.4.0-DMR 2023-10-12
6.5.5 2023-09-21
6.5.4 2023-08-28
7.3.0-DMR 2023-08-14
7.1.1 2023-07-24
6.1.7 2023-07-12
7.2.0-DMR 2023-06-29
6.5.3 2023-06-14
7.1.0 2023-05-31
6.5.2 2023-04-21
6.1.6 2023-04-12
7.0.0-DMR 2023-03-30
6.5.1 2023-03-10
6.1.5 2023-02-28
6.6.0-DMR 2023-02-20
6.1.4 2023-02-08
6.5.0 2022-12-29
5.1.5 2022-12-28
6.1.3 2022-12-05
5.3.4 2022-11-24
6.4.0-DMR 2022-11-17
6.1.2 2022-10-24
5.4.3 2022-10-13
6.3.0-DMR 2022-09-30
5.3.3 2022-09-14
6.1.1 2022-09-01
6.2.0-DMR 2022-08-23
5.4.2 2022-07-08
5.3.2 2022-06-29
6.1.0 2022-06-13
5.4.1 2022-05-13
5.2.4 2022-04-26
6.0.0-DMR 2022-04-07
3922
14.14.1.3.2 瀏覽器相容性
TiDB Dashboard 可在常見的、更新及時的桌面瀏覽器中使用，具體版本號為：
•Chrome >= 77
•Firefox >= 68
•Edge >= 17
注意：
若使用舊版本瀏覽器或其他瀏覽器訪問 TiDB Dashboard ，部分介面可能不能正常工作。
14.14.1.3.3 登入
訪問 TiDB Dashboard 將會顯示使用者登入介面。
•可使用 TiDB的root使用者登入。
•如果建立了 自定義 SQL使用者 ，也可以使用自定義的 SQL使用者和密碼登入。
3703

可以發現這裡面的資訊非常亂也很雜，大部分都是不相關的內容，依靠人工找出想要的資訊仍然是件費勁的事。但是藉助大模型的語義理解、邏輯推理、歸納總結等能力，我們就能得到一個非常清晰準確的答案。

如果把 TiDB 相關的各種文件、部落格、專欄文章、asktug問答等等內容全部放進 TiDB Vector 再加以調教，一個強大的問答機器人就誕生了。推薦閱讀：https://tidb.net/blog/a9cdb8ec

簡易版 TiDB Chat2Query

到這裡應該要結束了，但還有點意猶未盡。

用自然語言寫 SQL 是目前資料領域很熱門的一個方向，這個和本文主題並沒有太大關係，純好奇研究了一下，想給 TiDB 也做一個類似的工具。

其實 TiDB Cloud 很早就上線了 Chat2Query 功能：

企業微信截圖_20240426142647.png

比較容易想到的方案是把相關的表結構資訊和問題組合成 Prompt 一起發給大語言模型，類似於這樣：

企業微信截圖_20240509162700.png

落地到程式碼層面只需要簡單呼叫 OpenAI 的介面即可：

import os
from openai import OpenAI

os.environ['OPENAI_API_KEY'] = 'sk-xxx'
client = OpenAI()
completion = client.chat.completions.create(
  model="gpt-3.5-turbo",
  temperature=0,
  messages=[
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": """我的TiDB資料庫中有如下幾張表：
            Employee(id, name, department_id)
            Department(id, name, address)
            Salary_Payments(id, employee_id, amount, date)

            請幫我寫一段SQL查詢最近半年員工數超過10人的部門名稱。"""}
  ]
)
print(completion.choices[0].message)

如果每次問問題都要先把表結構查出來未免也太麻煩，有沒有辦法讓大模型一次性把所有表結構都記住，我們只問問題就行，最好是能根據問題查出最終資料？這個就涉及到 Agent（智慧體） 部分的內容了。

LangChain 提供了 SQL Agent 能力，只需要簡單幾行程式碼就可以把資料庫和大模型打通。

from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits import create_sql_agent
from langchain_openai import AzureChatOpenAI
import os

os.environ["OPENAI_API_TYPE"] = "azure"
os.environ["OPENAI_API_VERSION"] = "2024-02-01"
os.environ["AZURE_OPENAI_ENDPOINT"] = "https://heao-ai-test1.openai.azure.com/"
os.environ["OPENAI_API_KEY"] = "xxxxx"

db = SQLDatabase.from_uri("mysql+pymysql://root:123456@10.x.x.x:4000/test")

llm = AzureChatOpenAI(
        temperature=0.0,
        deployment_name="gpt-4",
        model_name="gpt-4" 
    )
agent_executor = create_sql_agent(llm, db=db, agent_type="openai-tools", verbose=True)
agent_executor.invoke("test庫下有多少張表")

企業微信截圖_20240514174658.png

再看稍微複雜點的例子：

insert into department values(1,'tidb team','F1'),(2,'mysql team','F2'),(3,'dev team','F3');
insert into employee values(1,'aaa',1),(2,'bbb',2),(3,'ccc',1),(4,'ddd',1),(5,'eee',3),(6,'fff',1);

agent_executor.invoke("查詢哪個部門的員工人數最多")

企業微信截圖_20240514181417.png

企業微信截圖_20240514181516.png

SELECT department.name, COUNT(employee.id) as employee_count
FROM department
JOIN employee ON department.id = employee.department_id
GROUP BY department.id
ORDER BY employee_count DESC
LIMIT 1

[('tidb team', 4)]The department with the most employees is the 'tidb team', which has 4 employees.

程式執行過程中先分析了和問題相關要使用的表，再根據語義生成了 SQL，最後把 SQL 傳送給 TiDB 得到最終結果，結果非常準確。

儘管看起來還不錯，但是實際使用中侷限性還比較大，比如：

  • 真實生產庫無法訪問外網大模型
  • 無法實現跨庫查詢
  • 無法使用資料庫本身的一些特性或讀取元資訊，比如想看show table regions這種
  • 複雜業務邏輯識別誤差較大，比較依賴提問技巧

之前也體驗過其他類似的產品，個人感覺現階段實用性還略有欠缺，但不妨礙它是 AI4DB 的一個熱門方向，相信未來會變得更好。

總結

藉助 TiDB 向量檢索能力，可以非常輕鬆地和 AI 生態進行打通，這也意味著 TiDB 的使用場景變得更加豐富。可以預見的是 AI 浪潮會持續火熱，可能以後向量檢索就成了資料庫的標配。前不久 Oracle 釋出了整合向量特性的新版本，直接更名為Oracle 23 ai炸響大半個資料庫圈子，DBA 們擁抱 AI 也必須要安排起來了。

作者介紹：hey-hoho，來自神州數碼鈦合金戰隊，是一支致力於為企業提供分散式資料庫TiDB整體解決方案的專業技術團隊。團隊成員擁有豐富的資料庫從業背景，全部擁有TiDB高階資格證書，並活躍於TiDB開源社群，是官方認證合作伙伴。目前已為10+客戶提供了專業的TiDB交付服務，涵蓋金融、證券、物流、電力、政府、零售等重點行業。

相關文章