使用 TiDB Vector 搭建 RAG 应用 - 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 的最新版本是多少，它的回答一定不让人满意。

如果我就想让 ChatGPT 告诉我 TiDB 的最新版本号，通常有两种办法：

• 全量模式：OpenAI 把 GPT 重新训练一次，生成新的 snapshot（超级烧钱）
• 增量模式：给 GPT 投喂一些新的信息，比如在问题的上下文中把新版本信息告诉 GPT，让它根据提供的上下文来回答

很明显第二种方法更切实际，工程实现上更容易，成本也更低。

我们从指定的文档、甚至是搜索引擎中找到相关的信息丢给 GPT ，借助 GPT 的推理能力并且限定它的上下文内容得到最终答案。通过临时喂给 LLM 的数据提升它的能力，这就构成了 RAG 的灵魂：检索（Retrieval）、增强（Augmented）、生成（Generation）。

这就是一个最基础的 RAG （也称之为朴素 RAG，Native RAG）流程，在此基础上如果继续优化提升准确度的话还可以引入 Rerank 等相关技术形成高级 RAG（Advance RAG）。

可以发现检索是 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 的一些元信息，如文件名、页号等。

文档预处理和 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 功能：

比较容易想到的方案是把相关的表结构信息和问题组合成 Prompt 一起发给大语言模型，类似于这样：

落地到代码层面只需要简单调用 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库下有多少张表")

再看稍微复杂点的例子：

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("查询哪个部门的员工人数最多")

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交付服务，涵盖金融、证券、物流、电力、政府、零售等重点行业。