Google AlloyDB for PostgreSQL - AlloyDBVectorStore

AlloyDB 是一项全托管关系型数据库服务，提供高性能、无缝集成和出色的可扩展性。AlloyDB 与 PostgreSQL 100% 兼容。扩展你的数据库应用，利用 AlloyDB 的 LlamaIndex 集成构建 AI 驱动的体验。

本 notebook 介绍如何使用 AlloyDB for PostgreSQL 通过 AlloyDBVectorStore 类存储向量嵌入。

在 GitHub 上详细了解该包。

开始之前

要运行此 notebook，你需要执行以下操作：

🦙 安装库

安装集成库 llama-index-alloydb-pg，以及嵌入服务库 llama-index-embeddings-vertex。

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

仅限 Colab： 取消注释以下单元格以重启内核，或使用按钮重启内核。对于 Vertex AI Workbench，你可以使用顶部按钮重启终端。

# # Automatically restart kernel after installs so that your environment can access the new packages
# import IPython

# app = IPython.Application.instance()
# app.kernel.do_shutdown(True)

🔐 身份验证

以登录此 notebook 的 IAM 用户身份向 Google Cloud 进行身份验证，以便访问你的 Google Cloud 项目。

如果你使用 Colab 运行此 notebook，请使用下面的单元格并继续。
如果你使用 Vertex AI Workbench，请查看此处的设置说明。

from google.colab import auth

auth.authenticate_user()

☁ 设置你的 Google Cloud 项目

设置你的 Google Cloud 项目，以便在此 notebook 中利用 Google Cloud 资源。

如果你不知道自己的项目 ID，可以尝试以下方法：

运行 gcloud config list。
运行 gcloud projects list。
查看支持页面：查找项目 ID。

# @markdown Please fill in the value below with your Google Cloud project ID and then run the cell.

PROJECT_ID = "my-project-id"  # @param {type:"string"}

# Set the project id
!gcloud config set project {PROJECT_ID}

基本用法

设置 AlloyDB 数据库值

在 AlloyDB 实例页面中查找你的数据库值。

# @title Set Your Values Here { display-mode: "form" }
REGION = "us-central1"  # @param {type: "string"}
CLUSTER = "my-cluster"  # @param {type: "string"}
INSTANCE = "my-primary"  # @param {type: "string"}
DATABASE = "my-database"  # @param {type: "string"}
TABLE_NAME = "vector_store"  # @param {type: "string"}
USER = "postgres"  # @param {type: "string"}
PASSWORD = "my-password"  # @param {type: "string"}

AlloyDBEngine 连接池

将 AlloyDB 设为向量存储的要求和参数之一是一个 AlloyDBEngine 对象。AlloyDBEngine 会为你的 AlloyDB 数据库配置连接池，使应用能够成功连接，并遵循行业最佳实践。

要使用 AlloyDBEngine.from_instance() 创建 AlloyDBEngine，你只需提供 5 项内容：

project_id：AlloyDB 实例所在 Google Cloud 项目的项目 ID。
region：AlloyDB 实例所在的区域。
cluster：AlloyDB 集群的名称。
instance：AlloyDB 实例的名称。
database：要在 AlloyDB 实例上连接的数据库名称。

默认情况下，将使用 IAM 数据库身份验证作为数据库身份验证方法。此库使用来自环境的应用默认凭据（ADC）所属的 IAM 主体。

此外，也可以使用内置数据库身份验证，通过用户名和密码访问 AlloyDB 数据库。只需向 AlloyDBEngine.from_instance() 提供可选的 user 和 password 参数：

user：用于内置数据库身份验证和登录的数据库用户
password：用于内置数据库身份验证和登录的数据库密码。

注意： 本教程演示异步接口。所有异步方法都有对应的同步方法。

from llama_index_alloydb_pg import AlloyDBEngine

engine = await AlloyDBEngine.afrom_instance(
    project_id=PROJECT_ID,
    region=REGION,
    cluster=CLUSTER,
    instance=INSTANCE,
    database=DATABASE,
    user=USER,
    password=PASSWORD,
)

用于 AlloyDB Omni 的 AlloyDBEngine

要为 AlloyDB Omni 创建 AlloyDBEngine，你需要一个连接 URL。AlloyDBEngine.from_connection_string 首先创建一个异步引擎，然后将其转换为 AlloyDBEngine。以下是使用 asyncpg 驱动的连接示例：

# Replace with your own AlloyDB Omni info
OMNI_USER = "my-omni-user"
OMNI_PASSWORD = ""
OMNI_HOST = "127.0.0.1"
OMNI_PORT = "5432"
OMNI_DATABASE = "my-omni-db"

connstring = f"postgresql+asyncpg://{OMNI_USER}:{OMNI_PASSWORD}@{OMNI_HOST}:{OMNI_PORT}/{OMNI_DATABASE}"
engine = AlloyDBEngine.from_connection_string(connstring)

初始化表

AlloyDBVectorStore 类需要一个数据库表。AlloyDBEngine 引擎有一个辅助方法 init_vector_store_table()，可用于为你创建具有正确 schema 的表。

await engine.ainit_vector_store_table(
    table_name=TABLE_NAME,
    vector_size=768,  # Vector size for VertexAI model(textembedding-gecko@latest)
)

可选提示：💡

你也可以在传入 table_name 的任何地方传入 schema_name 来指定 schema 名称。

SCHEMA_NAME = "my_schema"

await engine.ainit_vector_store_table(
    table_name=TABLE_NAME,
    schema_name=SCHEMA_NAME,
    vector_size=768,
)

创建嵌入类实例

你可以使用任何 Llama Index 嵌入模型。你可能需要启用 Vertex AI API 才能使用 VertexTextEmbeddings。我们建议在生产环境中设置嵌入模型的版本，详细了解文本嵌入模型。

# enable Vertex AI API
!gcloud services enable aiplatform.googleapis.com

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)

初始化默认的 AlloyDBVectorStore

from llama_index_alloydb_pg import AlloyDBVectorStore

vector_store = await AlloyDBVectorStore.create(
    engine=engine,
    table_name=TABLE_NAME,
    # schema_name=SCHEMA_NAME
)

下载数据

!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()
print("Document ID:", documents[0].doc_id)

与 VectorStoreIndex 一起使用

使用 VectorStoreIndex 从向量存储创建索引。

使用文档初始化向量存储

使用向量存储最简单的方法是加载一组文档，并使用 from_documents 从中构建索引。

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
)

查询索引

query_engine = index.as_query_engine()
response = query_engine.query("What did the author do?")
print(response)

创建自定义向量存储

向量存储可以利用关系数据来过滤相似度搜索。

创建一个带有自定义元数据列的新表。你也可以复用一个已经为 Document 的 id、content、embedding 和/或 metadata 设置了自定义列的现有表。

from llama_index_alloydb_pg import Column

# Set table name
TABLE_NAME = "vectorstore_custom"
# SCHEMA_NAME = "my_schema"

await engine.ainit_vector_store_table(
    table_name=TABLE_NAME,
    # schema_name=SCHEMA_NAME,
    vector_size=768,  # VertexAI model: textembedding-gecko@003
    metadata_columns=[Column("len", "INTEGER")],
)

# Initialize AlloyDBVectorStore
custom_store = await AlloyDBVectorStore.create(
    engine=engine,
    table_name=TABLE_NAME,
    # schema_name=SCHEMA_NAME,
    metadata_columns=["len"],
)

添加带元数据的文档

Document 的 metadata 可以为 LLM 和检索过程提供更多信息。详细了解提取和添加元数据的不同方法。

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
)

使用元数据过滤器搜索文档

你可以通过指定 filters 参数对搜索结果应用预过滤

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))

添加索引

通过应用向量索引来加速向量搜索查询。详细了解向量索引。

from llama_index_alloydb_pg.indexes import IVFFlatIndex

index = IVFFlatIndex()
await vector_store.aapply_vector_index(index)

ScaNN 索引创建（仅在 AlloyDB Omni 中可用）需要足够的维护工作内存。你需要在应用索引之前调用 set_maintenance_work_mem 设置数据库标志 maintenance_work_mem。

from llama_index_alloydb_pg.indexes import ScaNNIndex

VECTOR_SIZE = 768  # Replace with the vector size of your embedding model
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)

重建索引

await vector_store.areindex()  # Re-index using default index name

删除索引

await vector_store.adrop_vector_index()  # Delete index using default name

索引

数据加载

查询

存储

成本分析

应用

聊天机器人

问答

结构化数据

智能体

Chat Engines

Query Engine

连接器

文档与节点

摄取流水线

节点解析器

大语言模型

提示词

回调

节点后处理器

响应合成器

检索器

路由

结构化输出

现有数据

Presets and Modes

能力

起步

能力

Examples

Azure Deployment

Databases and Queues

LLM Integrations