Google Cloud SQL for PostgreSQL - `PostgresDocumentStore` 与 `PostgresIndexStore`

Cloud SQL is a fully managed relational database service that offers high performance, seamless integration, and impressive scalability. It offers MySQL, PostgreSQL, and SQL Server database engines. Extend your database application to build AI-powered experiences leveraging Cloud SQL’s LlamaIndex integrations.

本笔记本介绍如何使用 Cloud SQL for PostgreSQL 通过 PostgresDocumentStore 和 PostgresIndexStore 类来存储文档和索引。

Learn more about the package on GitHub.

开始之前

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

🦙 库安装

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

%pip install --upgrade --quiet llama-index-cloud-sql-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}

基本用法

设置 Cloud SQL 数据库值

Find your database values, in the Cloud SQL Instances page.

# @title Set Your Values Here { display-mode: "form" }
REGION = "us-central1"  # @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"}

PostgresEngine 连接池

将 Cloud SQL 设为向量存储的要求和参数之一是 PostgresEngine 对象。PostgresEngine 用于配置与 Cloud SQL 数据库的连接池，使应用程序能够成功连接并遵循行业最佳实践。

To create a PostgresEngine using PostgresEngine.from_instance() you need to provide only 4 things:

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

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 envionment.

有关IAM数据库认证的更多信息，请参阅：

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

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

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

from llama_index_cloud_sql_pg import PostgresEngine

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

初始化一个表格

The PostgresDocumentStore class requires a database table. The PostgresEngine 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,
)

初始化一个默认的 PostgresDocumentStore

from llama_index_cloud_sql_pg import PostgresDocumentStore

doc_store = await PostgresDocumentStore.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_cloud_sql_pg import PostgresIndexStore


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

index_store = await PostgresIndexStore.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
)