Google AlloyDB for PostgreSQL - `AlloyDBDocumentStore` & `AlloyDBIndexStore`

AlloyDB is a fully managed relational database service that offers high performance, seamless integration, and impressive scalability. AlloyDB is 100% compatible with PostgreSQL. Extend your database application to build AI-powered experiences leveraging AlloyDB’s LlamaIndex integrations.

本笔记本介绍如何使用 AlloyDB for PostgreSQL 通过 AlloyDBDocumentStore 和 AlloyDBIndexStore 类来存储文档和索引。

Learn more about the package on GitHub.

开始之前

要运行此笔记本，您需要执行以下操作：

🦙 库安装

Install the integration library, llama-index-alloydb-pg, and the library for the embedding service, llama-index-embeddings-vertex.

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

Colab only: Uncomment the following cell to restart the kernel or use the button to restart the kernel. For Vertex AI Workbench you can restart the terminal using the button on top.

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

🔐 身份验证

以登录此笔记本的IAM用户身份验证到Google Cloud，以便访问您的Google Cloud项目。

如果您正在使用 Colab 运行此笔记本，请使用下面的单元格并继续。
If you are using Vertex AI Workbench, check out the setup instructions here.

from google.colab import auth

auth.authenticate_user()

☁ 设置您的谷歌云项目

设置您的 Google Cloud 项目，以便在此笔记本中利用 Google Cloud 资源。

如果您不知道您的项目ID，请尝试以下方法：

Run gcloud config list.
Run gcloud projects list.
See the support page: Locate the project 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数据库值

Find your database values, in the AlloyDB Instances page.

# @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 = "document_store"  # @param {type: "string"}
USER = "postgres"  # @param {type: "string"}
PASSWORD = "my-password"  # @param {type: "string"}

AlloyDB引擎连接池

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

To create a AlloyDBEngine using AlloyDBEngine.from_instance() you need to provide only 5 things:

project_idproject_id : AlloyDB实例所在的Google Cloud项目的项目ID。
regionregion : AlloyDB实例所在的区域。
clustercluster: AlloyDB集群的名称。
instanceinstance : AlloyDB实例的名称。
databasedatabase : 要连接的 AlloyDB 实例上的数据库名称。

By default, IAM database authentication will be used as the method of database authentication. This library uses the IAM principal belonging to the Application Default Credentials (ADC) sourced from the environment.

Optionally, built-in database authentication using a username and password to access the AlloyDB database can also be used. Just provide the optional user and password arguments to AlloyDBEngine.from_instance():

useruser : 用于内置数据库认证和登录的数据库用户
passwordpassword : 用于内置数据库认证和登录的数据库密码。

Note: This tutorial demonstrates the async interface. All async methods have corresponding sync methods.

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)

初始化一个表格

The AlloyDBDocumentStore class requires a database table. The AlloyDBEngine engine has a helper method init_doc_store_table() that can be used to create a table with the proper schema for you.

await engine.ainit_doc_store_table(
    table_name=TABLE_NAME,
)

可选提示：💡

You can also specify a schema name by passing schema_name wherever you pass table_name.

SCHEMA_NAME = "my_schema"

await engine.ainit_doc_store_table(
    table_name=TABLE_NAME,
    schema_name=SCHEMA_NAME,
)

初始化一个默认的 AlloyDBDocumentStore

from llama_index_alloydb_pg import AlloyDBDocumentStore

doc_store = await AlloyDBDocumentStore.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)

解析为节点

from llama_index.core.node_parser import SentenceSplitter

nodes = SentenceSplitter().get_nodes_from_documents(documents)

设置索引存储

from llama_index_alloydb_pg import AlloyDBIndexStore


INDEX_TABLE_NAME = "index_store"
await engine.ainit_index_store_table(
    table_name=INDEX_TABLE_NAME,
)

index_store = await AlloyDBIndexStore.create(
    engine=engine,
    table_name=INDEX_TABLE_NAME,
    # schema_name=SCHEMA_NAME
)

添加到文档库

from llama_index.core import StorageContext

storage_context = StorageContext.from_defaults(
    docstore=doc_store, index_store=index_store
)

storage_context.docstore.add_documents(nodes)

与索引配合使用

文档存储库可与多个索引配合使用。每个索引使用相同的基础节点。

from llama_index.core import Settings, SimpleKeywordTableIndex, SummaryIndex
from llama_index.llms.vertex import Vertex

Settings.llm = Vertex(model="gemini-1.5-flash", project=PROJECT_ID)
summary_index = SummaryIndex(nodes, storage_context=storage_context)
keyword_table_index = SimpleKeywordTableIndex(
    nodes, storage_context=storage_context
)

查询索引

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

加载现有索引

文档存储库可与多个索引配合使用。每个索引使用相同的基础节点。

# note down index IDs
list_id = summary_index.index_id
keyword_id = keyword_table_index.index_id

from llama_index.core import load_index_from_storage

# re-create storage context
storage_context = StorageContext.from_defaults(
    docstore=doc_store, index_store=index_store
)

# load indices
summary_index = load_index_from_storage(
    storage_context=storage_context, index_id=list_id
)
keyword_table_index = load_index_from_storage(
    storage_context=storage_context, index_id=keyword_id
)