在适用于 RAG 的 LlamaIndex on Vertex AI 中使用 Vertex AI Feature Store

本页面介绍了如何将 Vertex AI Feature Store 设置为向量数据库,以便与适用于 RAG 的 LlamaIndex on Vertex AI 搭配使用。

适用于 RAG 的 LlamaIndex on Vertex AI 使用由 Spanner 提供支持的内置向量数据库来存储和管理文本文档的向量表示法。向量数据库会根据文档与给定查询的语义相似度检索相关文档。

通过将 Vertex AI Feature Store 集成为额外的向量数据库,适用于 RAG 的 LlamaIndex on Vertex AI 可以使用 Vertex AI Feature Store 以低延迟处理大量数据,从而有助于提高 RAG 应用的性能和可伸缩性。

设置 Vertex AI Feature Store

Vertex AI Feature Store 是一种托管式云原生服务,是 Vertex AI 的重要组件。它使您可以在 BigQuery 表或视图中管理特征数据,从而简化了机器学习 (ML) 特征管理和在线传送。这样便可实现低延迟的在线特征传送。

对于使用优化在线传送创建的 FeatureOnlineStore 实例,您可以利用向量相似度搜索来检索在语义上相似或相关的实体(称为近似最近邻)。

以下部分介绍了如何为 RAG 应用设置 Vertex AI Feature Store 实例。

创建 BigQuery 表架构

使用 Google Cloud 控制台创建 BigQuery 表架构。它必须包含以下字段才能用作数据源。

字段名称 数据类型 状态
corpus_id String 必需
file_id String 必需
chunk_id String 必需
chunk_data_type String 可以为 Null
chunk_data String 可以为 Null
file_original_uri String 可以为 Null
embeddings Float 重复

此代码示例演示了如何定义 BigQuery 表架构。

  # Use this sql query as reference for creating the table
  CREATE TABLE `your-project-id.input_us_central1.rag_source_new` (
    `corpus_id` STRING ,
    `file_id` STRING,
    `chunk_id` STRING,
    `chunk_data_type` STRING,
    `chunk_data` STRING,
    `embeddings` ARRAY<FLOAT64>,
    `file_original_uri` STRING,
  );

预配 FeatureOnlineStore 实例

如需实现特征在线传送,请使用 Vertex AI Feature Store CreateFeatureOnlineStore API 设置 FeatureOnlineStore 实例。如果您是首次预配 FeatureOnlineStore,此操作可能需要大约五分钟才能完成。

  # TODO(developer): Update and uncomment the following lines:
  # PROJECT_ID = "your-project-id"
  #
  # Set feature_online_store_id.
  # Example: "rag_fos_test"
  # FEATURE_ONLINE_STORE_ID="your-feature-online-store-id"

  # Call CreateFeatureOnlineStore to create a FeatureOnlineStore instance
  curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" https://us-central1-aiplatform.googleapis.com/v1/projects/${PROJECT_ID}/locations/us-central1/featureOnlineStores?feature_online_store_id=${FEATURE_ONLINE_STORE_ID}   -d '{
      "optimized": {},
  }'

  # TODO(developer): Update and uncomment the following lines:
  # Get operation_id returned in CreateFeatureOnlineStore
  # OPERATION_ID="your-operation-id"

  # Poll Operation status until done = true in the response
  curl -X GET \
  -H "Authorization: Bearer $(gcloud auth print-access-token)" \
  -H "Content-Type: application/json" \
  https://us-central1-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/us-central1/operations/${OPERATION_ID}

创建 FeatureView 资源

如需将存储特征数据源的 BigQuery 表连接到 FeatureOnlineStore 实例,请调用 CreateFeatureView API 以创建 FeatureView 资源。创建 FeatureView 资源时,请选择默认距离指标 DOT_PRODUCT_DISTANCE,该指标定义为点积的负数(DOT_PRODUCT_DISTANCE 越小表示相似度越高)。

此代码示例演示了如何创建 FeatureView 资源。

  # TODO(developer): Update and uncomment the following lines:
  # Set feature_view_id
  # Example: "feature_view_test"
  # FEATURE_VIEW_ID = "your-feature-view-id"
  #
  # The big_query_uri generated in the above BigQuery table schema creation step
  # The format should be "bq://" + BigQuery table ID
  # Example: "bq://tester.ragtest1.rag_testdata"
  # BIG_QUERY_URI=YOUR_BIG_QUERY_URI

  # Call CreateFeatureView API to create a FeatureView
  curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" \
  https://us-central1-aiplatform.googleapis.com/v1/projects/${PROJECT_ID}/locations/us-central1/featureOnlineStores/${FEATURE_ONLINE_STORE_ID}/featureViews?feature_view_id=${FEATURE_VIEW_ID} \
    -d '{
          "vertex_rag_source": {
            "uri": '\""${BIG_QUERY_URI}"\"'
          }
      }'

  # Call ListFeatureViews API to verify the FeatureView is created successfully
  curl -X GET -H "Authorization: Bearer $(gcloud auth application-default print-access-token)" -H "Content-Type: application/json" https://us-central1-aiplatform.googleapis.com/v1/projects/${PROJECT_ID}/locations/us-central1/featureOnlineStores/${FEATURE_ONLINE_STORE_ID}/featureViews

上传数据和在线传送

RAG API 负责处理数据上传和在线传送。

在适用于 RAG 的 LlamaIndex on Vertex AI 中使用 Vertex AI Feature Store

设置 Vertex AI Feature Store 实例后,以下部分将介绍如何将其设置为向量数据库以便与 RAG 应用搭配使用。

使用 Vertex AI Feature Store 实例作为向量数据库以创建 RAG 语料库

如需创建 RAG 语料库,您必须使用 FEATURE_VIEW_RESOURCE_NAME。RAG 语料库会进行创建并自动与 Vertex AI Feature Store 实例相关联。RAG API 使用生成的 rag_corpus_id 来处理到 Vertex AI Feature Store 实例的数据上传,并从 rag_corpus_id 检索相关上下文。

此代码示例演示了如何将 Vertex AI Feature Store 实例用作向量数据库以创建 RAG 语料库。

REST

# TODO(developer): Update and uncomment the following lines:
# CORPUS_DISPLAY_NAME = "your-corpus-display-name"
#
# Full feature view resource name
# Format: projects/${PROJECT_ID}/locations/us-central1/featureOnlineStores/${FEATURE_ONLINE_STORE_ID}/featureViews/${FEATURE_VIEW_ID}
# FEATURE_VIEW_RESOURCE_NAME = "your-feature-view-resource-name"

# Call CreateRagCorpus API to create a new RAG corpus
curl -X POST \
-H "Authorization: Bearer $(gcloud auth print-access-token)" \
-H "Content-Type: application/json" \
 https://us-central1-aiplatform.googleapis.com/v1beta1/projects//{PROJECT_ID}/locations/us-central1/ragCorpora -d '{
      "display_name" : '\""${CORPUS_DISPLAY_NAME}"\"',
      "rag_vector_db_config" : {
              "vertex_feature_store": {
                "feature_view_resource_name":'\""${FEATURE_VIEW_RESOURCE_NAME}"\"'
              }
        }
  }'

# Call ListRagCorpora API to verify the RAG corpus is created successfully
curl -sS -X GET \
-H "Content-Type: application/json" \
-H "Authorization: Bearer $(gcloud auth print-access-token)" \
"https://us-central1-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/us-central1/ragCorpora"

Python

import vertexai
from vertexai.preview import rag
from vertexai.preview.generative_models import GenerativeModel, Tool
# Set Project
PROJECT_ID = "YOUR_PROJECT_ID"  # @param {type:"string"}
vertexai.init(project=PROJECT_ID, location="us-central1")

# Configure a Google first-party embedding model
embedding_model_config = rag.EmbeddingModelConfig(
    publisher_model="text-embedding-004"
)

# Configure a Vertex AI Feature Store instance for the corpus
FEATURE_VIEW_RESOURCE_NAME = "YOUR_FEATURE_VIEW_RESOURCE_NAME"  # @param {type:"string"}
vector_db = rag.VertexFeatureStore(
   resource_name=FEATURE_VIEW_RESOURCE_NAME,
)

# Name your corpus
DISPLAY_NAME = "YOUR_DISPLAY_NAME"  # @param {type:"string"}

rag_corpus = rag.create_corpus(
    display_name=DISPLAY_NAME, embedding_model_config=embedding_model_config, vector_db=vector_db
)

# Check the corpus just created
rag.list_corpora()

使用 RAG API 将文件导入 BigQuery 表

使用 ImportRagFiles API 将文件从 Google Cloud Storage 或 Google 云端硬盘导入 Vertex AI Feature Store 实例的 BigQuery 表中。这些文件会嵌入到 BigQuery 表中并存储在其中。

此代码示例演示了如何使用 RAG API 将文件导入 BigQuery 表。

REST

# TODO(developer): Update and uncomment the following lines:
# RAG_CORPUS_ID = "your-rag-corpus-id"
#
# Google Cloud Storage bucket/file location.
# For example, "gs://rag-fos-test/"
# GCS_URIS= "your-gcs-uris"

# Call ImportRagFiles API to embed files and store in the BigQuery table
curl -X POST \
-H "Authorization: Bearer $(gcloud auth print-access-token)" \
-H "Content-Type: application/json" \
https://us-central1-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/us-central1/ragCorpora/${RAG_CORPUS_ID}/ragFiles:import \
-d '{
  "import_rag_files_config": {
    "gcs_source": {
      "uris": '\""${GCS_URIS}"\"'
    },
    "rag_file_chunking_config": {
      "chunk_size": 512
    }
  }
}'

# Call ListRagFiles API to verify the files are imported successfully
curl -X GET \
-H "Content-Type: application/json" \
-H "Authorization: Bearer $(gcloud auth print-access-token)" \
https://us-central1-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/us-central1/ragCorpora/${RAG_CORPUS_ID}/ragFiles

Python

RAG_CORPUS_RESOURCE = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/YOUR_RAG_CORPUS_ID"
GS_BUCKET = "YOUR_GS_BUCKET"

response = rag.import_files(
    corpus_name=RAG_CORPUS_RESOURCE,
    paths=[GS_BUCKET],
    chunk_size=512,  # Optional
    chunk_overlap=100,  # Optional
)

运行同步过程以构建 FeatureOnlineStore 索引

将数据上传到 BigQuery 表后,运行同步过程以便使数据可进行在线传送。您必须使用 FeatureView 生成 FeatureOnlineStore 索引,并且同步过程可能需要 20 分钟才能完成。

此代码示例演示了如何运行同步过程以构建 FeatureOnlineStore 索引。

  # Call Feature Store SyncFeatureView API to run the synchronization process
  curl   "https://us-central1-aiplatform.googleapis.com/v1/${FEATURE_VIEW_RESOURCE_NAME}:sync" \
    -X POST \
    -H "Authorization: Bearer $(gcloud auth print-access-token)" \
    -H "Content-Type: application/json; charset=utf-8"

  # TODO(developer): Update and uncomment the following lines:
  # Call Vertex AI Feature Store GetFeatureViewSync API to check the running synchronization // status
  # FEATURE_VIEW_SYNC_ID = "your-feature-view-sync-id" returned in SyncFeatureView
  curl   "https://us-central1-aiplatform.googleapis.com/v1/${FEATURE_VIEW_RESOURCE_NAME}/featureViewSyncs/${FEATURE_VIEW_SYNC_ID}" \
    -X GET \
    -H "Authorization: Bearer $(gcloud auth print-access-token)" \
    -H "Content-Type: application/json; charset=utf-8"

使用 RAG API 检索相关上下文

同步过程完成后,您可以通过 RetrieveContexts API 从 FeatureOnlineStore 索引检索相关上下文。

REST

# TODO(developer): Update and uncomment the following lines:
# RETRIEVAL_QUERY="your-retrieval-query"
#
# Full RAG corpus resource name
# Format:
# "projects/${PROJECT_ID}/locations/us-central1/ragCorpora/${RAG_CORPUS_ID}"
# RAG_CORPUS_RESOURCE="your-rag-corpus-resource"

# Call RetrieveContexts API to retrieve relevant contexts
curl -X POST \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer $(gcloud auth print-access-token)" \
https://us-central1-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/us-central1:retrieveContexts \
  -d '{
    "vertex_rag_store": {
      "rag_resources": {
          "rag_corpus": '\""${RAG_CORPUS_RESOURCE}"\"',
        },
    },
    "query": {
      "text": '\""${RETRIEVAL_QUERY}"\"',
      "similarity_top_k": 10
    }
  }'
 ```

Python

RAG_CORPUS_RESOURCE = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/YOUR_RAG_CORPUS_ID"
RETRIEVAL_QUERY = "YOUR_RETRIEVAL_QUERY"

response = rag.retrieval_query(
    rag_resources=[
        rag.RagResource(
             rag_corpus=RAG_CORPUS_RESOURCE,
             # Optional: supply IDs from `rag.list_files()`.
             # rag_file_ids=["rag-file-1", "rag-file-2", ...],
        )
    ],
    text=RETRIEVAL_QUERY,
    similarity_top_k=10,  # Optional
)
print(response)

使用 Vertex AI Gemini API 生成内容

调用 Vertex AI GenerateContent API 以使用 Gemini 模型生成内容,并在请求中指定 RAG_CORPUS_RESOURCE 以从 FeatureOnlineStore 索引检索数据。

REST

# TODO(developer): Update and uncomment the following lines:
# MODEL_ID=gemini-1.5-flash-001
# GENERATE_CONTENT_PROMPT="your-generate-content-prompt"

# GenerateContent with contexts retrieved from the FeatureStoreOnline index
curl -X POST \
-H "Authorization: Bearer $(gcloud auth print-access-token)" \
-H "Content-Type: application/json"  https://us-central1-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/us-central1/publishers/google/models/${MODEL_ID}:generateContent \
-d '{
  "contents": {
    "role": "user",
    "parts": {
      "text": '\""${GENERATE_CONTENT_PROMPT}"\"'
    }
  },
  "tools": {
    "retrieval": {
      "vertex_rag_store": {
        "rag_resources": {
            "rag_corpus": '\""${RAG_CORPUS_RESOURCE}"\"',
          },
        "similarity_top_k": 8,
      }
    }
  }
}'

Python

RAG_CORPUS_RESOURCE = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/YOUR_RAG_CORPUS_ID"

rag_resource = rag.RagResource(
    rag_corpus=RAG_CORPUS_RESOURCE,
    # Optional: supply IDs from `rag.list_files()`.
    # rag_file_ids=["rag-file-1", "rag-file-2", ...],
)

rag_retrieval_tool = Tool.from_retrieval(
    retrieval=rag.Retrieval(
        source=rag.VertexRagStore(
            rag_resources=[
                rag.RagResource(
                    rag_corpus=RAG_CORPUS_RESOURCE,
                    # Optional: supply IDs from `rag.list_files()`.
                    # rag_file_ids=["rag-file-1", "rag-file-2", ...],
                )
            ],
            similarity_top_k=10,  # Optional
        ),
    )
)

rag_model = GenerativeModel(
   model_name="gemini-1.5-flash-001", tools=[rag_retrieval_tool]
)

GENERATE_CONTENT_PROMPT="YOUR_GENERATE_CONTENT_PROMPT"

response = rag_model.generate_content(GENERATE_CONTENT_PROMPT)
print(response.text)

后续步骤