将向量搜索集成到 LlamaIndex

本教程演示如何将 TiDB 的 vector search 功能与 LlamaIndex 集成。

前提条件

完成本教程，你需要：

• 安装 Python 3.8 或更高版本
• 安装 Jupyter Notebook
• 安装 Git
• 一个 TiDB 集群

如果你还没有 TiDB 集群，可以按照以下步骤创建：

• （推荐）参考 创建 TiDB Cloud Starter 集群 来创建你自己的 TiDB Cloud 集群。
• 参考 部署本地测试 TiDB 集群 或 部署生产环境 TiDB 集群 来创建版本为 v8.4.0 或更高的本地集群。

开始使用

本节提供逐步指导，帮助你将 TiDB 向量搜索与 LlamaIndex 集成，实现语义搜索。

步骤 1. 创建新的 Jupyter Notebook 文件

在根目录下，创建一个名为 integrate_with_llamaindex.ipynb 的 Jupyter Notebook 文件：

touch integrate_with_llamaindex.ipynb

步骤 2. 安装所需依赖

在你的项目目录下，运行以下命令安装所需的包：

pip install llama-index-vector-stores-tidbvector
pip install llama-index

在 Jupyter Notebook 中打开 integrate_with_llamaindex.ipynb 文件，并添加以下代码导入所需包：

import textwrap

from llama_index.core import SimpleDirectoryReader, StorageContext
from llama_index.core import VectorStoreIndex
from llama_index.vector_stores.tidbvector import TiDBVectorStore

步骤 3. 配置环境变量

根据你选择的 TiDB 部署方式，配置环境变量。

# 使用 getpass 在终端中安全提示输入环境变量。
import getpass
import os

# 连接字符串格式："mysql+pymysql://<USER>:<PASSWORD>@<HOST>:4000/<DB>?ssl_ca=/etc/ssl/cert.pem&ssl_verify_cert=true&ssl_verify_identity=true"
tidb_connection_string = getpass.getpass("TiDB Connection String:")
os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API Key:")

TIDB_DATABASE_URL="mysql+pymysql://<USERNAME>:<PASSWORD>@<HOST>:<PORT>/<DATABASE_NAME>"
# 例如：TIDB_DATABASE_URL="mysql+pymysql://root@127.0.0.1:4000/test"

步骤 4. 加载示例文档

步骤 4.1 下载示例文档

在你的项目目录下，创建一个名为 data/paul_graham/ 的目录，并从 run-llama/llama_index GitHub 仓库下载示例文档 paul_graham_essay.txt。

!mkdir -p 'data/paul_graham/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'

步骤 4.2 加载文档

使用 SimpleDirectoryReader 类加载 data/paul_graham/paul_graham_essay.txt 中的示例文档。

documents = SimpleDirectoryReader("./data/paul_graham").load_data()
print("Document ID:", documents[0].doc_id)

for index, document in enumerate(documents):
   document.metadata = {"book": "paul_graham"}

步骤 5. 嵌入并存储文档向量

步骤 5.1 初始化 TiDB 向量存储

以下代码在 TiDB 中创建一个名为 paul_graham_test 的表，用于优化向量搜索。

tidbvec = TiDBVectorStore(
   connection_string=tidb_connection_url,
   table_name="paul_graham_test",
   distance_strategy="cosine",
   vector_dimension=1536,
   drop_existing_table=False,
)

成功执行后，你可以直接在 TiDB 数据库中查看并访问 paul_graham_test 表。

步骤 5.2 生成并存储嵌入

以下代码解析文档，生成嵌入，并存入 TiDB 向量存储中。

storage_context = StorageContext.from_defaults(vector_store=tidbvec)
index = VectorStoreIndex.from_documents(
   documents, storage_context=storage_context, show_progress=True
)

预期输出如下：

Parsing nodes: 100%|██████████| 1/1 [00:00<00:00,  8.76it/s]
Generating embeddings: 100%|██████████| 21/21 [00:02<00:00,  8.22it/s]

步骤 6. 执行向量搜索

以下代码基于 TiDB 向量存储创建查询引擎，并执行语义相似度搜索。

query_engine = index.as_query_engine()
response = query_engine.query("What did the author do?")
print(textwrap.fill(str(response), 100))

预期输出如下：

The author worked on writing, programming, building microcomputers, giving talks at conferences,
publishing essays online, developing spam filters, painting, hosting dinner parties, and purchasing
a building for office use.

步骤 7. 使用元数据过滤器进行搜索

你可以使用元数据过滤器，细化搜索结果，获取符合过滤条件的最近邻结果。

使用 book != "paul_graham" 过滤器进行查询

以下示例排除 book 元数据字段为 "paul_graham" 的结果：

from llama_index.core.vector_stores.types import (
   MetadataFilter,
   MetadataFilters,
)

query_engine = index.as_query_engine(
   filters=MetadataFilters(
      filters=[
         MetadataFilter(key="book", value="paul_graham", operator="!="),
      ]
   ),
   similarity_top_k=2,
)
response = query_engine.query("What did the author learn?")
print(textwrap.fill(str(response), 100))

预期输出如下：

Empty Response

使用 book == "paul_graham" 过滤器进行查询

以下示例过滤结果，只包含 book 元数据字段为 "paul_graham" 的文档：

from llama_index.core.vector_stores.types import (
   MetadataFilter,
   MetadataFilters,
)

query_engine = index.as_query_engine(
   filters=MetadataFilters(
      filters=[
         MetadataFilter(key="book", value="paul_graham", operator="=="),
      ]
   ),
   similarity_top_k=2,
)
response = query_engine.query("What did the author learn?")
print(textwrap.fill(str(response), 100))

预期输出如下：

The author learned programming on an IBM 1401 using an early version of Fortran in 9th grade, then
later transitioned to working with microcomputers like the TRS-80 and Apple II. Additionally, the
author studied philosophy in college but found it unfulfilling, leading to a switch to studying AI.
Later on, the author attended art school in both the US and Italy, where they observed a lack of
substantial teaching in the painting department.

步骤 8. 删除文档

删除索引中的第一个文档：

tidbvec.delete(documents[0].doc_id)

检查文档是否已被删除：

query_engine = index.as_query_engine()
response = query_engine.query("What did the author learn?")
print(textwrap.fill(str(response), 100))

预期输出如下：

Empty Response

相关链接

• Vector Data Types
• Vector Search Index