【RAG】【vector_stores002】Google AlloyDB for PostgreSQL 向量存储完整案例

本案例演示如何使用 LlamaIndex 与 Google AlloyDB for PostgreSQL 集成，实现向量存储和检索功能，用于构建基于文档的问答系统。

1. 案例目标

本案例的主要目标是：

1. 设置 AlloyDB 向量存储：配置 LlamaIndex 以使用 Google AlloyDB for PostgreSQL 作为向量数据库。
2. 文档索引与存储：将文档内容加载并存储到 AlloyDB 向量数据库中。
3. 查询与检索：基于存储的文档内容，实现自然语言查询并获取相关答案。
4. 元数据过滤：演示如何使用元数据过滤来精确控制搜索结果。
5. 向量索引优化：展示如何添加向量索引以提高查询性能。

2. 技术栈与核心依赖

• LlamaIndex：用于构建基于 LLM 的应用程序的框架
• llama-index-alloydb-pg：LlamaIndex 的 AlloyDB 向量存储插件
• llama-index-embeddings-vertex：Vertex AI 文本嵌入模型
• llama-index-llms-vertex：Vertex AI 大语言模型
• Google AlloyDB for PostgreSQL：Google Cloud 提供的完全托管的关系型数据库服务

核心依赖安装

%pip install --upgrade --quiet llama-index-alloydb-pg llama-index-embeddings-vertex llama-index-llms-vertex llama-index

3. 环境配置

在使用本案例前，需要完成以下环境配置：

1. Google Cloud 项目：创建一个 Google Cloud 项目
2. 启用 AlloyDB API：在 Google Cloud 控制台中启用 AlloyDB API
3. AlloyDB 集群和实例：创建 AlloyDB 集群和实例
4. AlloyDB 数据库：创建 AlloyDB 数据库
5. 数据库用户：为数据库添加用户
6. 身份验证：配置 Google Cloud 身份验证

4. 案例实现

4.1 身份验证和项目设置

# 身份验证 from google.colab import auth auth.authenticate_user() # 设置 Google Cloud 项目 PROJECT_ID = "my-project-id" # 替换为您的项目ID !gcloud config set project {PROJECT_ID}

4.2 设置 AlloyDB 数据库参数

# AlloyDB 数据库连接参数 REGION = "us-central1" # 替换为您的区域 CLUSTER = "my-cluster" # 替换为您的集群名称 INSTANCE = "my-primary" # 替换为您的实例名称 DATABASE = "my-database" # 替换为您的数据库名称 TABLE_NAME = "vector_store" # 替换为您的表名 USER = "postgres" # 替换为您的用户名 PASSWORD = "my-password" # 替换为您的密码

4.3 创建 AlloyDB 引擎

from llama_index_alloydb_pg import AlloyDBEngine # 使用实例信息创建 AlloyDB 引擎 engine = await AlloyDBEngine.afrom_instance( project_id=PROJECT_ID, region=REGION, cluster=CLUSTER, instance=INSTANCE, database=DATABASE, user=USER, password=PASSWORD, )

4.4 初始化向量存储表

# 初始化向量存储表 await engine.ainit_vector_store_table( table_name=TABLE_NAME, vector_size=768, # VertexAI 模型的向量大小 )

4.5 配置嵌入模型和 LLM

# 启用 Vertex AI API !gcloud services enable aiplatform.googleapis.com # 配置嵌入模型和 LLM from llama_index.core import Settings from llama_index.embeddings.vertex import VertexTextEmbedding from llama_index.llms.vertex import Vertex import google.auth credentials, project_id = google.auth.default() Settings.embed_model = VertexTextEmbedding( model_name="textembedding-gecko@003", project=PROJECT_ID, credentials=credentials, ) Settings.llm = Vertex(model="gemini-1.5-flash-002", project=PROJECT_ID)

4.6 创建向量存储

from llama_index_alloydb_pg import AlloyDBVectorStore # 创建向量存储 vector_store = await AlloyDBVectorStore.create( engine=engine, table_name=TABLE_NAME, )

4.7 准备数据并创建索引

# 下载数据 !mkdir -p 'data/paul_graham/' !wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt' # 加载文档 from llama_index.core import SimpleDirectoryReader documents = SimpleDirectoryReader("./data/paul_graham").load_data() # 创建索引 from llama_index.core import StorageContext, VectorStoreIndex storage_context = StorageContext.from_defaults(vector_store=vector_store) index = VectorStoreIndex.from_documents( documents, storage_context=storage_context, show_progress=True )

4.8 查询索引

# 创建查询引擎并执行查询 query_engine = index.as_query_engine() response = query_engine.query("What did the author do?") print(response)

4.9 创建自定义向量存储

from llama_index_alloydb_pg import Column # 设置表名 TABLE_NAME = "vectorstore_custom" # 创建带有自定义元数据列的表 await engine.ainit_vector_store_table( table_name=TABLE_NAME, vector_size=768, # VertexAI 模型: textembedding-gecko@003 metadata_columns=[Column("len", "INTEGER")], ) # 初始化自定义向量存储 custom_store = await AlloyDBVectorStore.create( engine=engine, table_name=TABLE_NAME, metadata_columns=["len"], )

4.10 添加带元数据的文档

from llama_index.core import Document # 创建带元数据的文档 fruits = ["apple", "pear", "orange", "strawberry", "banana", "kiwi"] documents = [ Document(text=fruit, metadata={"len": len(fruit)}) for fruit in fruits ] storage_context = StorageContext.from_defaults(vector_store=custom_store) custom_doc_index = VectorStoreIndex.from_documents( documents, storage_context=storage_context, show_progress=True )

4.11 使用元数据过滤搜索

from llama_index.core.vector_stores.types import ( MetadataFilter, MetadataFilters, FilterOperator, ) # 创建过滤器 filters = MetadataFilters( filters=[ MetadataFilter(key="len", operator=FilterOperator.GT, value="5"), ], ) # 使用过滤器查询 query_engine = custom_doc_index.as_query_engine(filters=filters) res = query_engine.query("List some fruits") print(str(res.source_nodes[0].text))

4.12 添加向量索引

from llama_index_alloydb_pg.indexes import IVFFlatIndex # 添加 IVFFlat 索引 index = IVFFlatIndex() await vector_store.aapply_vector_index(index)

4.13 添加 ScaNN 索引（仅适用于 AlloyDB Omni）

from llama_index_alloydb_pg.indexes import ScaNNIndex # 设置向量大小 VECTOR_SIZE = 768 # 替换为您的嵌入模型的向量大小 # 创建 ScaNN 索引 index = ScaNNIndex(name="my_scann_index") await vector_store.aset_maintenance_work_mem(index.num_leaves, VECTOR_SIZE) await vector_store.aapply_vector_index(index)

4.14 重新索引和删除索引

# 重新索引 await vector_store.areindex() # 使用默认索引名称重新索引 # 删除索引 await vector_store.adrop_vector_index() # 使用默认名称删除索引

5. 案例效果

通过本案例，您可以实现：

• 将文档内容向量化并存储到 Google AlloyDB for PostgreSQL 中
• 使用自然语言查询文档内容，如 "What did the author do?"
• 获取基于文档内容的准确回答
• 使用元数据过滤精确控制搜索结果
• 通过添加向量索引提高查询性能

6. 案例实现思路

本案例的核心实现思路是：

1. 环境准备：配置 Google Cloud 项目和 AlloyDB 实例
2. 数据库连接：使用AlloyDBEngine建立与 AlloyDB 的连接
3. 表初始化：创建适合向量存储的数据库表结构
4. 模型配置：配置 Vertex AI 的嵌入模型和语言模型
5. 向量存储：使用AlloyDBVectorStore创建向量存储
6. 文档处理：使用SimpleDirectoryReader加载文档
7. 索引创建：使用VectorStoreIndex将文档内容向量化并存储
8. 查询执行：通过query_engine执行自然语言查询
9. 高级功能：实现元数据过滤和向量索引优化

7. 扩展建议

• 批量文档处理：扩展系统以支持批量处理多个文档
• 高级元数据过滤：实现更复杂的元数据过滤条件，如范围查询、多条件组合等
• 混合搜索：结合向量搜索和传统文本搜索提高检索精度
• 性能优化：针对大规模文档集优化索引和查询性能
• 实时更新：实现文档内容的实时更新和增量索引
• 多语言支持：扩展系统以支持多语言文档处理和查询
• 自定义嵌入模型：集成其他嵌入模型替代 Vertex AI
• 分布式处理：利用 AlloyDB 的分布式特性处理大规模数据

8. 总结

本案例展示了如何使用 LlamaIndex 与 Google AlloyDB for PostgreSQL 集成，构建一个基于向量存储的文档问答系统。通过将文档内容向量化并存储在 AlloyDB 中，我们可以实现高效的语义搜索和准确的问答功能。AlloyDB 作为 Google Cloud 提供的完全托管的关系型数据库服务，提供了高性能、高可用性和可扩展性，特别适合需要处理大量文档并提供智能问答的企业应用场景。