主题
本 notebook 展示如何使用基于 Postgres 的向量存储 NileVectorStore,为多租户 RAG 应用存储和查询向量嵌入。
什么是 Nile?
Nile 是一个 Postgres 数据库,可为每个租户启用所有数据库操作,包括自动扩展、分支和备份,并实现完全的客户隔离。
多租户 RAG 应用越来越受欢迎,因为它们在利用大语言模型的同时提供了安全性和隐私性。
然而,管理底层 Postgres 数据库并不简单。每租户一个数据库(DB-per-tenant)成本高昂且管理复杂,而共享数据库(shared-DB)存在安全和隐私问题,还限制了 RAG 应用的可扩展性和性能。Nile 对 Postgres 进行了重新设计,以兼得两者之长——在共享数据库的成本、效率和开发者体验下,实现每租户一个数据库的隔离性。
在共享数据库中存储数百万个向量可能很慢,并且需要大量资源来建立索引和查询。但如果你将 1000 个租户存储在 Nile 的虚拟租户数据库中,每个租户有 1000 个向量,这将是相当可控的。特别是因为你可以将较大的租户放在它们自己的计算资源上,而较小的租户可以高效地共享计算资源并按需自动扩展。
Nile 入门
首先注册 Nile。注册 Nile 后,系统会提示你创建第一个数据库。继续操作。你将被重定向到新数据库的 “Query Editor” 页面。
在那里,点击 “Home”(左侧菜单顶部的图标),点击 “generate credentials” 并复制生成的连接字符串。你很快就会用到它。
其他资源
开始之前
安装依赖
让我们安装并导入依赖项。
如果你在 Colab 中打开这个 Notebook,可能需要先安装 LlamaIndex 🦙。
%pip install llama-index-vector-stores-nile
%pip install /Users/gwen/workspaces/llama_index/llama-index-integrations/vector_stores/llama-index-vector-stores-nile/dist/llama_index_vector_stores_nile-0.1.1.tar.gz!pip install llama-indeximport logging
from llama_index.core import SimpleDirectoryReader, StorageContext
from llama_index.core import VectorStoreIndex
from llama_index.core.vector_stores import (
MetadataFilter,
MetadataFilters,
FilterOperator,
)
from llama_index.vector_stores.nile import NileVectorStore, IndexType设置与 Nile 数据库的连接
假设你已按照上一节“Nile 入门”中的说明操作,你现在应该已获得 Nile 数据库的连接字符串。
你可以将其设置在名为 NILEDB_SERVICE_URL 的环境变量中,或直接在 Python 中设置。
%env NILEDB_SERVICE_URL=postgresql://username:password@us-west-2.db.thenile.dev:5432/niledb现在,我们将创建一个 NileVectorStore。请注意,除了 URL 和维度等常规参数外,我们还设置了 tenant_aware=True。
🔥 NileVectorStore 同时支持租户感知向量存储(为每个租户隔离文档)和常规存储(通常用于所有租户都可访问的共享数据)。下面,我们将演示租户感知向量存储。
# Get the service url by reading local .env file with NILE_SERVICE_URL variable
import os
NILEDB_SERVICE_URL = os.environ["NILEDB_SERVICE_URL"]
# OR set it explicitly
# NILE_SERVICE_URL = "postgresql://nile:password@db.thenile.dev:5432/nile"
vector_store = NileVectorStore(
service_url=NILEDB_SERVICE_URL,
table_name="documents",
tenant_aware=True,
num_dimensions=1536,
)设置 OpenAI
你可以将其设置在 .env 文件中,或直接在 Python 中设置
%env OPENAI_API_KEY=sk-...# Uncomment and set it explicitly if you prefer not to use .env
# os.environ["OPENAI_API_KEY"] = "sk-..."多租户相似度搜索
为了演示使用 LlamaIndex 和 Nile 进行多租户相似度搜索,我们将下载两份文档——每份都包含来自不同公司销售电话的转录文本。Nexiv 提供 IT 服务,ModaMart 从事零售业务。我们将为每份文档添加租户标识符,并将它们加载到租户感知向量存储中。然后,我们将为每个租户查询该存储。你将看到同一个问题如何生成两个不同的响应,因为它为每个租户检索到了不同的文档。
下载数据
!mkdir -p data
!wget "https://raw.githubusercontent.com/niledatabase/niledatabase/main/examples/ai/sales_insight/data/transcripts/nexiv-solutions__0_transcript.txt" -O "data/nexiv-solutions__0_transcript.txt"
!wget "https://raw.githubusercontent.com/niledatabase/niledatabase/main/examples/ai/sales_insight/data/transcripts/modamart__0_transcript.txt" -O "data/modamart__0_transcript.txt"加载文档
我们将使用 LlamaIndex 的 SimpleDirectoryReader 加载文档。由于我们希望在加载后使用租户元数据更新文档,因此将为每个租户使用单独的读取器
reader = SimpleDirectoryReader(
input_files=["data/nexiv-solutions__0_transcript.txt"]
)
documents_nexiv = reader.load_data()
reader = SimpleDirectoryReader(input_files=["data/modamart__0_transcript.txt"])
documents_modamart = reader.load_data()使用租户元数据丰富文档
我们将创建两个 Nile 租户,并将每个租户的租户 ID 添加到文档元数据中。我们还会添加一些额外的元数据,例如自定义文档 ID 和类别。此元数据可用于在检索过程中过滤文档。当然,在你自己的应用中,你也可以为现有租户加载文档,并添加你认为有用的任何元数据信息。
tenant_id_nexiv = str(vector_store.create_tenant("nexiv-solutions"))
tenant_id_modamart = str(vector_store.create_tenant("modamart"))
# Add the tenant id to the metadata
for i, doc in enumerate(documents_nexiv, start=1):
doc.metadata["tenant_id"] = tenant_id_nexiv
doc.metadata[
"category"
] = "IT" # We will use this to apply additional filters in a later example
doc.id_ = f"nexiv_doc_id_{i}" # We are also setting a custom id, this is optional but can be useful
for i, doc in enumerate(documents_modamart, start=1):
doc.metadata["tenant_id"] = tenant_id_modamart
doc.metadata["category"] = "Retail"
doc.id_ = f"modamart_doc_id_{i}"使用 NileVectorStore 创建 VectorStore 索引
我们将所有文档加载到同一个 VectorStoreIndex 中。由于我们在设置时创建了租户感知的 NileVectorStore,Nile 将正确使用元数据中的 tenant_id 字段来隔离它们。
将没有 tenant_id 的文档加载到租户感知存储中会抛出 ValueException。
storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex.from_documents(
documents_nexiv + documents_modamart,
storage_context=storage_context,
show_progress=True,
)/Users/gwen/.pyenv/versions/3.10.15/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
from .autonotebook import tqdm as notebook_tqdm
Parsing nodes: 100%|██████████| 2/2 [00:00<00:00, 1129.32it/s]
Generating embeddings: 100%|██████████| 2/2 [00:00<00:00, 4.58it/s]为每个租户查询索引
你可以在下面看到我们如何为每个查询指定租户,因此我们得到的答案与该租户相关,并且仅属于他们
nexiv_query_engine = index.as_query_engine(
similarity_top_k=3,
vector_store_kwargs={
"tenant_id": str(tenant_id_nexiv),
},
)
print(nexiv_query_engine.query("What were the customer pain points?"))The customer pain points were related to managing customer data using multiple platforms, leading to data discrepancies, time-consuming reconciliation efforts, and decreased productivity.modamart_query_engine = index.as_query_engine(
similarity_top_k=3,
vector_store_kwargs={
"tenant_id": str(tenant_id_modamart),
},
)
print(modamart_query_engine.query("What were the customer pain points?"))The customer's pain points were concerns about the quality and value of the winter jackets, skepticism towards reviews, worries about sizing and fit when ordering clothes online, and the desire for a warm but lightweight jacket.查询现有嵌入
在上面的示例中,我们通过加载和嵌入新文档来创建索引。但如果我们已经生成了嵌入并将它们存储在 Nile 中呢?在这种情况下,你仍然像上面那样初始化 NileVectorStore,但不要使用 VectorStoreIndex.from_documents(...),而是使用以下代码:
index = VectorStoreIndex.from_vector_store(vector_store=vector_store)
query_engine = index.as_query_engine(
vector_store_kwargs={
"tenant_id": str(tenant_id_modamart),
},
)
response = query_engine.query("What action items do we need to follow up on?")
print(response)The action items to follow up on include sending the customer detailed testimonials about the lightweight and warm qualities of the jackets, providing the customer with a sizing guide, and emailing the customer a 10% discount on their first purchase.使用 ANN 索引进行近似最近邻搜索
Nile 支持 pgvector 支持的所有索引——IVFFlat 和 HNSW。IVFFlat 更快、使用资源更少且易于调优。HNSW 创建和使用时消耗更多资源,调优更具挑战性,但在精度/速度权衡方面表现出色。让我们看看如何使用索引,尽管一个包含两份文档的示例实际上并不需要它们。
IVFFlat 索引
IVFFlat 索引的工作原理是将向量空间划分为称为 “lists” 的区域,首先找到最近的列表,然后在这些列表中搜索最近邻。你在创建索引时指定列表数量(nlists),然后在查询时可以指定搜索中将使用多少个最近列表(ivfflat_probes)。
try:
vector_store.create_index(index_type=IndexType.PGVECTOR_IVFFLAT, nlists=10)
except Exception as e:
# This will throw an error if the index already exists, which may be expected
print(e)
nexiv_query_engine = index.as_query_engine(
similarity_top_k=3,
vector_store_kwargs={
"tenant_id": str(tenant_id_nexiv),
"ivfflat_probes": 10,
},
)
print(
nexiv_query_engine.query("What action items do we need to follow up on?")
)
vector_store.drop_index()Index documents_embedding_idx already existsHNSW 索引
HNSW 索引的工作原理是将向量空间划分为多层图,其中每一层包含不同粒度级别上点之间的连接。在搜索期间,它从粗粒度层导航到更细粒度层,识别数据中的最近邻。在创建索引期间,你指定一层中的最大连接数(m)以及构建图时考虑的候选向量数量(ef_construction)。查询时,你可以指定将要搜索的候选列表大小(hnsw_ef)。
try:
vector_store.create_index(
index_type=IndexType.PGVECTOR_HNSW, m=16, ef_construction=64
)
except Exception as e:
# This will throw an error if the index already exists, which may be expected
print(e)
nexiv_query_engine = index.as_query_engine(
similarity_top_k=3,
vector_store_kwargs={
"tenant_id": str(tenant_id_nexiv),
"hnsw_ef": 10,
},
)
print(nexiv_query_engine.query("Did we mention any pricing?"))
vector_store.drop_index()其他 VectorStore 操作
元数据过滤器
NileVectorStore 还支持基于元数据过滤向量。例如,当我们加载文档时,为每份文档都包含了 category 元数据。现在我们可以使用此信息来过滤检索到的文档。请注意,此过滤是在租户过滤器之外额外进行的。在租户感知向量存储中,租户过滤器是强制性的,以防止意外的数据泄漏。
filters = MetadataFilters(
filters=[
MetadataFilter(
key="category", operator=FilterOperator.EQ, value="Retail"
),
]
)
nexiv_query_engine_filtered = index.as_query_engine(
similarity_top_k=3,
filters=filters,
vector_store_kwargs={"tenant_id": str(tenant_id_nexiv)},
)
print(
"test query on nexiv with filter on category = Retail (should return empty): ",
nexiv_query_engine_filtered.query("What were the customer pain points?"),
)test query on nexiv with filter on category = Retail (should return empty): Empty Response删除文档
删除文档可能非常重要。特别是当你的一些租户位于需要遵守 GDPR 的地区时。
ref_doc_id = "nexiv_doc_id_1"
vector_store.delete(ref_doc_id, tenant_id=tenant_id_nexiv)
# Query the data again
print(
"test query on nexiv after deletion (should return empty): ",
nexiv_query_engine.query("What were the customer pain points?"),
)test query on nexiv after deletion (should return empty): Empty Response