RAG Engine で Vertex AI ベクトル検索を使用する

このページでは、RAG Engine を Vertex AI ベクトル検索に接続する方法について説明します。

RAG Engine は、Spanner を搭載した組み込みベクトル データベースを使用して、テキスト ドキュメントのベクトル表現の保存と管理を行う強力なツールです。ベクトル データベースを使用すると、指定されたクエリに対するドキュメントの意味的な類似性に基づいて、関連するドキュメントを効率的に取得できます。Vertex AI ベクトル検索を追加のベクトル データベースとして RAG Engine と統合することで、ベクトル検索の機能を使用して低レイテンシで大量のデータを処理し、RAG アプリケーションのパフォーマンスとスケーラビリティを向上させることができます。

Vertex AI ベクトル検索の設定

Vertex AI ベクトル検索は、Google Research が開発したベクトル検索技術をベースにしています。ベクトル検索では、Google 検索、YouTube、Google Play などの Google プロダクトの基盤と同じインフラストラクチャを利用できます。

RAG Engine と統合するには、空のベクトル検索インデックスが必要です。

Vertex AI SDK を設定する

RAG アプリケーション用の Vertex AI Vector Search インスタンスを準備する手順は次のとおりです。

  1. Vertex AI SDK を設定するには、設定をご覧ください。

  2. 環境変数を次のように設定します。

    PROJECT_ID=YOUR_PROJECT_ID
    LOCATION=YOUR_LOCATION_ID
    
  3. 省略可: Vertex AI Workbench を使用している場合は、事前に認証されているため、この手順は必要ありません。それ以外の場合は、ノートブックを実行するために、次のセル認証を実行する必要があります。

    # If it's Colab runtime, authenticate the user with Google Cloud
    if "google.colab" in sys.modules:
        from google.colab import auth
    
        auth.authenticate_user()
    
  4. 次のコマンドを入力して API を有効にします。

    ! gcloud services enable compute.googleapis.com aiplatform.googleapis.com --project "{PROJECT_ID}"

aiplatform SDK を初期化する

aiplatform SDK を初期化するには、次の操作を行います。

# init the aiplatform package
from google.cloud import aiplatform
aiplatform.init(project=PROJECT_ID, location=LOCATION)

ベクトル検索インデックスを作成する

RAG コーパスと互換性のあるベクトル検索インデックスを作成するには、インデックスが次の条件を満たしている必要があります。

  1. IndexUpdateMethodSTREAM_UPDATE にする必要があります。ストリーム インデックスを作成するをご覧ください。

  2. 距離の測定タイプは、次のいずれかに明示的に設定する必要があります。

    • DOT_PRODUCT_DISTANCE
    • COSINE_DISTANCE
  3. ベクトルのディメンションは、RAG コーパスで使用するエンベディング モデルと一致している必要があります。その他のパラメータは、選択内容に基づいてチューニングできます。選択内容によって、追加のパラメータをチューニングできるかどうかが決まります。

# create the index
my_index = aiplatform.MatchingEngineIndex.create_tree_ah_index(
    display_name="your-display-name",
    description="your-description",
    dimensions=768,
    approximate_neighbors_count=10,
    leaf_node_embedding_count=500,
    leaf_nodes_to_search_percent=7,
    distance_measure_type="DOT_PRODUCT_DISTANCE",
    feature_norm_type="UNIT_L2_NORM",
    index_update_method="STREAM_UPDATE",
)

ベクトル検索インデックス エンドポイントを作成する

パブリック エンドポイントは RAG Engine でサポートされています。

# create IndexEndpoint
my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint.create(
    display_name="your-display-name", public_endpoint_enabled=True
)

インデックスをインデックス エンドポイントにデプロイする

最近傍検索を行う前に、インデックスをインデックス エンドポイントにデプロイする必要があります。

DEPLOYED_INDEX_ID="YOUR_DEPLOYED_INDEX_ID"

my_index_endpoint.deploy_index(index=my_index, deployed_index_id=DEPLOYED_INDEX_ID)

インデックスをインデックス エンドポイントに初めてデプロイする場合は、バックエンドを自動的にビルドして起動するまでに 30 分ほどかかります。その後、インデックスを保存できます。最初のデプロイ後、インデックスは数秒で準備が整います。インデックスのデプロイ ステータスを確認するには、ベクトル検索コンソールを開き、[インデックス エンドポイント] タブを選択して、インデックス エンドポイントを選択します。

インデックスとインデックス エンドポイントのリソース名を特定します。形式は次のとおりです。

  • projects/${PROJECT_ID}/locations/${LOCATION_ID}/indexes/${INDEX_ID}
  • projects/${PROJECT_ID}/locations/${LOCATION_ID}/indexEndpoints/${INDEX_ENDPOINT_ID}

リソース名がわからない場合は、次のコマンドを使用して確認します。

print(my_index_endpoint.resource_name)
print(my_index.resource_name)

RAG Engine で Vertex AI ベクトル検索を使用する

Vector Search インスタンスが設定されたら、このセクションの手順に沿って、ベクトル検索インスタンスを RAG アプリケーションのベクトル データベースとして設定します。

ベクトル データベースを設定して RAG コーパスを作成する

RAG コーパスを作成する場合は、完全な INDEX_ENDPOINT_NAMEINDEX_NAME のみを指定します。RAG コーパスが作成され、ベクトル検索インデックスに自動的に関連付けられます。検証は条件に対して実行されます。いずれかの要件が満たされていない場合、リクエストは拒否されます。

Python

CORPUS_DISPLAY_NAME = "YOUR_CORPUS_DISPLAY_NAME"
index_resource_name = my_index.resource_name
endpoint_resource_name = my_index_endpoint.resource_name
vector_db = rag.VertexVectorSearch(index=index_resource_name, index_endpoint=endpoint_resource_name)
rag_corpus = rag.create_corpus(display_name=CORPUS_DISPLAY_NAME, vector_db=vector_db)

REST

# TODO(developer): Update and un-comment the following lines:
# CORPUS_DISPLAY_NAME = "YOUR_CORPUS_DISPLAY_NAME"
# Full index/indexEndpoint resource name
# Index: projects/${PROJECT_ID}/locations/${LOCATION_ID}/indexes/${INDEX_ID}
# IndexEndpoint: projects/${PROJECT_ID}/locations/${LOCATION_ID}/indexEndpoints/${INDEX_ENDPOINT_ID}
# INDEX_RESOURCE_NAME = "YOUR_INDEX_ENDPOINT_RESOURCE_NAME"
# INDEX_NAME = "YOUR_INDEX_RESOURCE_NAME"
# Call CreateRagCorpus API to create a new RagCorpus
curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" https://${LOCATION_ID}-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/${LOCATION_ID}/ragCorpora -d '{
      "display_name" : '\""${CORPUS_DISPLAY_NAME}"\"',
      "rag_vector_db_config" : {
              "vertex_vector_search": {
                "index":'\""${INDEX_NAME}"\"'
            "index_endpoint":'\""${INDEX_ENDPOINT_NAME}"\"'
              }
        }
  }'

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

省略可: ベクトル検索の情報なしで RAG コーパスを作成する

後で更新できるベクトル検索情報のない空の RAG コーパスを作成するには、次のいずれかのコードサンプルを選択します。

Python

CORPUS_DISPLAY_NAME = "YOUR_CORPUS_DISPLAY_NAME"
vector_db = rag.VertexVectorSearch()
rag_corpus = rag.create_corpus(display_name=CORPUS_DISPLAY_NAME, vector_db=vector_db)

REST

# TODO(developer): Update and un-comment the following lines:
# Call CreateRagCorpus API to create a new RAG corpus without the Vector Search information.
curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" https://${LOCATION_ID}-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/${LOCATION_ID}/ragCorpora -d '{
      "display_name" : '\""${CORPUS_DISPLAY_NAME}"\"',
      "rag_vector_db_config" : {
              "vertex_vector_search": {}
        }
  }'

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

ベクトル検索リソースを設定したら、対応する情報で RAG コーパスを更新できます。

Python

index_resource_name = my_index.resource_name
endpoint_resource_name = my_index_endpoint.resource_name
vector_db = rag.VertexVectorSearch(index=index_resource_name, index_endpoint=endpoint_resource_name)
updated_rag_corpus = rag.update_corpus(corpus_name=rag_corpus.name, vector_db=vector_db)

REST

curl -X PATCH \
-H "Authorization: Bearer $(gcloud auth print-access-token)" \
-H "Content-Type: application/json" \
https://${LOCATION_ID}-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/${LOCATION_ID}/ragCorpora -d '{
      "rag_vector_db_config" : {
"vertex_vector_search": {
                "index":'\""${INDEX_NAME}"\"'
                "index_endpoint":'\""${INDEX_ENDPOINT_NAME}"\"'
              }
}
  }'

RAG API を使用してファイルをインポートする

ImportRagFiles API を使用して、Cloud Storage または Google ドライブからベクトル検索インデックスにファイルをインポートします。ファイルは埋め込まれ、ベクトル検索インデックスに保存されます。

Python

RAG_CORPUS_RESOURCE = "projects/{PROJECT_ID}/locations/{LOCATION_ID}/ragCorpora/YOUR_RAG_CORPUS_ID"
GCS_BUCKET = "YOUR_GCS_BUCKET"

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

REST

# TODO(developer): Update and un-comment the following lines:
# RAG_CORPUS_ID = "YOUR_RAG_CORPUS_ID"
//
# Google Cloud Storage bucket and 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://${LOCATION_ID}-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/${LOCATION_ID}/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 that the files are imported successfully
curl -X GET \
-H "Content-Type: application/json" \
-H "Authorization: Bearer $(gcloud auth print-access-token)" \
https://${LOCATION_ID}-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/${LOCATION_ID}/ragCorpora/${RAG_CORPUS_ID}/ragFiles

RAG API を使用して関連するコンテキストを取得する

ファイルのインポートが完了したら、RetrieveContexts API を使用してベクトル検索インデックスから関連するコンテキストを取得できます。

Python

RAG_CORPUS_RESOURCE = "projects/{PROJECT_ID}/locations/{LOCATION_ID}/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
    vector_distance_threshold=0.3,  # Optional
)
print(response)

REST

# TODO(developer): Update and un-comment the following lines:
# RETRIEVAL_QUERY="YOUR_RETRIEVAL_QUERY"

# Full RagCorpus resource name
# Format:
# "projects/${PROJECT_ID}/locations/${LOCATION_ID}/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://${LOCATION_ID}-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/${LOCATION_ID}:retrieveContexts \
  -d '{
    "vertex_rag_store": {
      "rag_resources": {
          "rag_corpus": '\""${RAG_CORPUS_RESOURCE}"\"',
        },
      "vector_distance_threshold": 0.3
    },
    "query": {
      "text": '\""${RETRIEVAL_QUERY}"\"',
      "similarity_top_k": 10
    }
  }'

Vertex AI Gemini API を使用してコンテンツを生成する

Gemini モデルを使用してコンテンツを生成する場合は、Vertex AI GenerateContent API を呼び出します。リクエストで RAG_CORPUS_RESOURCE を指定すると、API はベクトル検索インデックスからデータを自動的に取得します。

Python

from vertexai.preview.generative_models import GenerativeModel, Tool

RAG_CORPUS_RESOURCE = "projects/{PROJECT_ID}/locations/{LOCATION_ID}/ragCorpora/YOUR_RAG_CORPUS_ID"

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
            vector_distance_threshold=0.3,   # 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)

REST

# TODO(developer): Update and un-comment the following lines:
# MODEL_ID=gemini-pro
# GENERATE_CONTENT_PROMPT="YOUR_GENERATE_CONTENT_PROMPT"

# GenerateContent with contexts retrieved from the Vector Search index

curl -X POST \
-H "Authorization: Bearer $(gcloud auth print-access-token)" \
-H "Content-Type: application/json"  https://${LOCATION_ID}-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/${LOCATION_ID}/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,
        "vector_distance_threshold": 0.32
      }
    }
  }
}'

次のステップ