RAG Engine で Vertex AI Feature Store を使用する

プレビュー

RAG Engine はプレビュー版で、Google Cloud サービス固有の規約の「pre-GA サービス規約」が適用されます。pre-GA のプロダクトと機能では、サポートが制限される場合があり、他の pre-GA のバージョンと互換性のない変更が行われる可能性があります。詳細については、リリース ステージの説明をご覧ください。さらに、Vertex AI で Gemini API を使用すると、生成 AI プレビュー版の利用規約（プレビュー規約）に同意したことになります。

このページでは、RAG Engine で使用するベクトル データベースとして Vertex AI Feature Store を設定する方法について説明します。

RAG Engine は、Spanner を搭載した組み込みベクトル データベースを使用して、テキスト ドキュメントのベクトル表現の保存と管理を行います。ベクトル データベースは、指定されたクエリに対するドキュメントの意味的な類似性に基づいて、関連するドキュメントを取得します。

Vertex AI Feature Store を追加のベクトル データベースとして統合することで、RAG Engine は低レイテンシで大量のデータを処理できるため、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 を初めてプロビジョニングする場合、オペレーションが完了するまでに 5 分ほどかかることがあります。

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

次のステップ

• グラウンディングの詳細を確認する。グラウンディングの概要をご覧ください。
• RAG Engine の詳細を確認する。RAG Engine を使用するをご覧ください。
• グラウンディングと RAG の詳細を確認する。RAG を使用したレスポンスのグラウンディングをご覧ください。