Use Vector Search 2.0 with RAG

This page shows you how Vertex AI RAG Engine uses Vertex AI Vector Search 2.0, which is a Google Cloud product that serves as a comprehensive storage and retrieval system to store and manage vector representations of your documents. The Vertex AI Vector Search 2.0 is then used to retrieve relevant documents based on the document's semantic similarity to a given query.

RAG will manage the Vector Search 2.0 collections for you and you will have full access to the Vector Search 2.0 collections in your project.

Choosing a RAG backend

When using Vertex AI RAG Engine, you can choose between different backend storage options for your RAG corpora. The following table summarizes the key differences:

Backend Option Underlying Technology Management by RAG Engine Data Visibility in Project CMEK Support
RagManagedVertexVectorSearch Vertex AI Vector Search 2.0 Fully managed by Google. Visible in your Google Cloud project No
VertexVectorSearch Vertex AI Vector Search 1.0 User is expected to set up, manage, and clean up the Vector Search 1.0 instance. Visible in your Google Cloud project No
RagManagedDb Spanner Fully managed by Google. Not directly visible in your project Yes

Choose RagManagedVertexVectorSearch for the latest features, automated management, and transparency into your Vector Search 2.0 data. If CMEK support is a strict requirement, consider RagManagedDb, keeping in mind the lack of direct visibility into the backend database.

Enable Vector Search API

To use RAG Managed Vertex AI Vector Search, you need to first enable the Vector Search API.

Create a RAG corpus with RagManagedVertexVectorSearch

This code samples demonstrates how to create a RAG corpus using RagManagedVertexVectorSearch.

Python 

from vertexai.preview import rag
import vertexai

PROJECT_ID = YOUR_PROJECT_ID
LOCATION = YOUR_RAG_ENGINE_LOCATION
DISPLAY_NAME = YOUR_RAG_CORPUS_DISPLAY_NAME

# Initialize Vertex AI API once per session
vertexai.init(project=PROJECT_ID, location=LOCATION)

vector_db = rag.RagManagedVertexVectorSearch()
rag_corpus = rag.create_corpus(
    display_name=DISPLAY_NAME, backend_config=rag.RagVectorDbConfig(vector_db=vector_db))

REST

Replace the following variables:

  • PROJECT_ID: Your project ID.
  • LOCATION: The region to process the request.
  • CORPUS_DISPLAY_NAME: The display name of the RAG corpus.
PROJECT_ID=PROJECT_ID
LOCATION=LOCATION
CORPUS_DISPLAY_NAME=CORPUS_DISPLAY_NAME

curl -X POST \
-H "Authorization: Bearer $(gcloud auth print-access-token)" \
-H "Content-Type: application/json" \
https://${LOCATION}-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/${LOCATION}/ragCorpora \
-d '{
      "display_name" : '\""${CORPUS_DISPLAY_NAME}"\"',
      "vector_db_config": {
        "rag_managed_vertex_vector_search": {}
      }
    }'

Importing your data into RagManagedVertexVectorSearch

You can use either the ImportRagFiles API or the UploadRagFile API to import your data into the RagManagedVertexVectorSearch.

To upload your local file into your RAG corpus, see Upload a RAG file. To import data into your RAG corpus, see the following code sample that demonstrates how to import from Cloud Storage. To learn about the supported data sources, see Data sources supported for RAG.

Python 

from vertexai.preview import rag
import vertexai

PROJECT_ID = YOUR_PROJECT_ID
LOCATION = YOUR_RAG_ENGINE_LOCATION
CORPUS_ID = YOUR_CORPUS_ID
PATHS = ["gs://my_bucket/my_files_dir"]

# Initialize Vertex AI API once per session
vertexai.init(project=PROJECT_ID, location=LOCATION)

corpus_name = f"projects/{PROJECT_ID}/locations/{LOCATION}/ragCorpora/{CORPUS_ID}"
# This is a non blocking call.
response = await rag.import_files_async(
    corpus_name=corpus_name,
    paths=PATHS,
)

# Wait for the import to complete.
await response.result()

REST 

GCS_URI=GCS_URI

curl -X POST \
-H "Authorization: Bearer $(gcloud auth print-access-token)" \
-H "Content-Type: application/json" \
https://${LOCATION}-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/${LOCATION}/ragCorpora/${CORPUS_ID}/ragFiles:import \
-d '{
  "import_rag_files_config": {
    "gcs_source": {
      "uris": '\""${GCS_URI}"\"',
      },
  }
}'

What's next