Use a Weaviate database with RAG Engine

This page shows you how to connect your RAG Engine corpus to your Weaviate database.

You can use your Weaviate database instance, which is an open source database, with RAG Engine to index and conduct a vector-based similarity search. A similarity search is a way to find pieces of text that are similar to the text that you're looking for, which requires the use of an embedding model. The embedding model produces vector data for each piece of text being compared. The similarity search is used to retrieve semantic contexts for grounding to return the most accurate content from your LLM.

With RAG Engine, you can continue to use your fully-managed vector database instance, which you are responsible for provisioning. RAG Engine uses the vector database for storage, index management, and search.

Considerations

Consider the following steps before using the Weaviate database:

1. You must create, configure, and deploy your Weaviate database instance and collection. Follow the instructions in Create your Weaviate collection to set up a collection based on your schema.
2. You must provide a Weaviate API key, which allows RAG Engine to interact with the Weaviate database. RAG Engine supports the API key-based AuthN and AuthZ, which connects to your Weaviate database and supports an HTTPS connection.
3. RAG Engine doesn't store and manage your Weaviate API key. Instead, you must do the following:
   1. Store your key in the Google Cloud Secret Manager.
   2. Grant your project's service account permissions to access your secret.
   3. Provide RAG Engine access to your secret's resource name.
   4. When you interact with your Weaviate database, RAG Engine accesses your secret resource using your service account.
4. RAG Engine corpus and the Weaviate collection have a one-to-one mapping. RAG files are stored in a Weaviate database collection. When a call is made to the CreateRagCorpus API or the UpdateRagCorpus API, the RAG corpus is associated to the database collection.
5. In addition to dense embeddings-based semantic searches, the hybrid search is also supported with RAG Engine through a Weaviate database. You can also adjust the weight between dense and sparse vector similarity in a hybrid search.

Provision the Weaviate database

Before using the Weaviate database with RAG Engine, you must do the following:

1. Configure and deploy your Weaviate database instance.
2. Prepare the HTTPS endpoint.
3. Create your Weaviate collection.
4. Use your API key to provision Weaviate using AuthN and AuthZ.
5. Provision your RAG Engine service account.

Configure and deploy your Weaviate database instance

You must follow the Weaviate official guide quickstart. However, you can use the Google Cloud Marketplace guide, which is optional.

You can set up your Weaviate instance anywhere as long as the Weaviate endpoint is accessible to configure and deploy in your project. You can then fully manage your Weaviate database instance.

Because RAG Engine isn't involved in any stage of your Weaviate database instance lifecycle, it is your responsibility to grant permissions to RAG Engine so it can store and search for data in your Weaviate database. It is also your responsibility to ensure that the data in your database can be used by RAG Engine. For example, if you change your data, RAG Engine isn't responsible for any unexpected behaviors because of those changes.

Prepare the HTTPS endpoint

During Weaviate provisioning, ensure that you create an HTTPS endpoint. Although HTTP connections are supported, we prefer that RAG Engine and Weaviate database traffic use an HTTPS connection.

Create your Weaviate collection

Because the RAG Engine corpus and the Weaviate collection have a one-to-one mapping, you must create a collection in your Weaviate database before associating your collection with the RAG Engine corpus. This one-time association is made when you call the CreateRagCorpus API or the UpdateRagCorpus API.

When creating a collection in Weaviate, you must use the following schema:

Use your API key to provision Weaviate using AuthN and AuthZ

Provisioning the Weaviate API key involves the following steps:

1. Create the Weaviate API key.
2. Configure Weaviate using your Weaviate API key.
3. Store your Weaviate API key in Secret Manager.

Create the API key

RAG Engine can only connect to your Weaviate database instances by using your API key for authentication and authorization. You must follow the Weaviate official guide to authentication to configure the API key-based authentication in your Weaviate database instance.

If creating the Weaviate API key requires identity information to associate with that comes from RAG Engine, you must create your first corpus, and use your RAG Engine service account as an identity.

Store your API key in Secret Manager

An API key holds Sensitive Personally Identifiable Information (SPII), which is subject to legal requirements. If the SPII data is compromised or misused, an individual might experience a significant risk or harm. To minimize risks to an individual while using RAG Engine, don't store and manage your API key, and avoid sharing the unencrypted API key.

To protect SPII, do the following:

1. Store your API key in Secret Manager.
2. Grant your RAG Engine service account the permissions to your secret(s), and manage the access control at the secret resource level.
   1. Navigate to your project's permissions.
   2. Enable the option Include Google-provided role grants.
   3. Find the service account, which has the format

      service-{project number}@gcp-sa-vertex-rag.iam.gserviceaccount.com

   4. Edit the service account's principals.
   5. Add the Secret Manager Secret Accessor role to the service account.
3. During the creation or update of the RAG corpus, pass the secret resource name to RAG Engine, and store the secret resource name.

When you make API requests to your Weaviate database instance(s), RAG Engine uses each service account to read the API key that corresponds to your secret resources in Secret Manager from your project(s).

Provision your RAG Engine service account

When you create the first resource in your project, RAG Engine creates a dedicated service account. You can find your service account from your project's IAM page. The service account follows this format:

service-{project number}@gcp-sa-vertex-rag.iam.gserviceaccount.com

For example, service-123456789@gcp-sa-vertex-rag.iam.gserviceaccount.com.

When integrating with the Weaviate database, your service account is used in the following scenarios:

• You can use your service account to generate your Weaviate API key for authentication. In some cases, generating the API key doesn't require any user information, which means that a service account isn't required when generating the API key.
• You can bind your service account with the API key in your Weaviate database to configure the authentication (AuthN) and authorization (AuthZ). However, your service account isn't required.
• You can store the API key Secret Manager in your project, and you can grant your service account permissions to these secret resources.
• RAG Engine uses service accounts to access the API key from the Secret Manager in your projects.

Set up your Google Cloud console environment

Prepare your RAG corpus

To access data from your Weaviate database, RAG Engine must have access to a RAG corpus. This section provides the steps for creating a single RAG corpus and additional RAG corpora.

Use CreateRagCorpus and UpdateRagCorpus APIs

You must specify the following fields when calling the CreateRagCorpus and UpdateRagCorpus APIs:

• rag_vector_db_config.weaviate: After you call the CreateRagCorpus API, the vector database configuration is chosen. The vector database configuration contains all of the configuration fields. If the rag_vector_db_config.weaviate field isn't set, then rag_vector_db_config.rag_managed_db is set by default.
• weaviate.http_endpoint: The HTTPS or HTTP Weaviate endpoint is created during provisioning of the Weaviate database instance.
• weaviate.collection_name: The name of the collection that is created during the Weaviate instance provisioning. The name must start with a capital letter.
• api_auth.api_key_config: The configuration specifies to use an API key to authorize your access to the vector database.
• api_key_config.api_key_secret_version: The resource name of the secret that is stored in Secret Manager, which contains your Weaviate API key.

You can create and associate your RAG corpus to the Weaviate collection in your database instance. However, you might need the service account to generate your API key and to configure your Weaviate database instance. When you create your first RAG corpus, the service account is generated. After you create your first RAG corpus, the association between the Weaviate database and the API key might not be ready for use in the creation of another RAG corpus.

Just in case your database and key aren't ready to be associated to your RAG corpus, do the following to your RAG corpus:

1. Set the weaviate field in rag_vector_db_config.

   • You can't change the associated vector database.
   • Leave both the http_endpoint and the collection_name fields empty. Both fields can be updated at a later time.

2. If you don't have your API key stored in Secret Manager, then you can leave the api_auth field empty. When you call the UpdateRagCorpus API, you can update the api_auth field. Weaviate requires that the following be done:

   1. Set the api_key_config in the api_auth field.

   2. Set the api_key_secret_version of your Weaviate API key in Secret Manager. The api_key_secret_version field uses the following format:

      projects/{project}/secrets/{secret}/versions/{version}

3. If you specify fields that can only be set one time, like http_endpoint or collection_name, you can't change them unless you delete your RAG corpus, and create your RAG corpus again. Other fields like the API key field, api_key_secret_version, can be updated.

4. When you call UpdateRagCorpus, you can set the vector_db field. The vector_db should be set to weaviate by your CreateRagCorpus API call. Otherwise, the system chooses the RAG Managed Database option, which is the default. This option can't be changed when you call the UpdateRagCorpus API. When you call UpdateRagCorpus and the vector_db field is partially set, you can update the fields that are marked as Changeable (also referred to as mutable).

This table lists the WeaviateConfig mutable and immutable fields that are used in your code.

Create the first RAG corpus

When the RAG Engine service account doesn't exist, do the following:

1. Create a RAG corpus in RAG Engine with an empty Weaviate configuration, which initiates RAG Engine provisioning to create a service account.
2. Choose a name for your RAG Engine service account that follows this format:
   

   service-{project number}@gcp-sa-vertex-rag.iam.gserviceaccount.com

   For example, service-123456789@gcp-sa-vertex-rag.iam.gserviceaccount.com.

3. Using your service account, access your secret that is stored in your project's Secret Manager, which contains your Weaviate API key.
4. Get the following information after Weaviate provisioning completes:
   • Your Weaviate HTTPS or HTTP endpoint.
   • The name of your Weaviate collection.
5. Call the CreateRagCorpus API to create a RAG corpus with an empty Weaviate configuration, and call the UpdateRagCorpus API to update the RAG corpus with the following information:
   • Your Weaviate HTTPS or HTTP endpoint.
   • The name of your Weaviate collection.
   • The API key resource name.

Create another RAG corpus

When the RAG Engine service account exists, do the following:

1. Get your RAG Engine service account from your project's permissions.
2. Enable the option "Include Google-provided role grants"
3. Choose a name for your RAG Engine service account that follows this format:
   

   service-{project number}@gcp-sa-vertex-rag.iam.gserviceaccount.com

4. Using your service account, access your secret that is stored in your project's Secret Manager, which contains your Weaviate API key.
5. During Weaviate provisioning, get the following information:
   • The Weaviate HTTPS or HTTP endpoint.
   • The name of your Weaviate collection.
6. Create a RAG corpus in RAG Engine, and connect with your Weaviate collection by doing one of the following:
   1. Make a CreateRagCorpus API call to create a RAG corpus with a populated Weaviate configuration, which is the preferred option.
   2. Make a CreateRagCorpus API call to create a RAG corpus with an empty Weaviate configuration, and make an UpdateRagCorpus API call to update the RAG corpus with the following information:
      • Weaviate database HTTP endpoint
      • Weaviate Collection name
      • API key

Examples

This section presents sample code that demonstrates how to set up your Weaviate database, Secret Manager, the RAG corpus, and the RAG file. Sample code is also provided to demonstrate how to import files, to retrieve context, to generate content, and to delete the RAG corpus and RAG files.

To use the Model Garden RAG API notebook, see Use Weaviate with Llama 3.

Set up your Weaviate database

This code sample demonstrates how to set up your Weaviate data and the Secret Manager.

# TODO(developer): Update the variables.
# The HTTPS/HTTP Weaviate endpoint you created during provisioning.
HTTP_ENDPOINT_NAME="https://your.weaviate.endpoint.com"

# Your Weaviate API Key.
WEAVIATE_API_KEY="example-api-key"

# Select your Weaviate collection name, which roughly corresponds to a Vertex AI Knowledge Engine Corpus.
# For example, "MyCollectionName"
# Note that the first letter needs to be capitalized.
# Otherwise, Weavaite will capitalize it for you.
WEAVIATE_COLLECTION_NAME="MyCollectionName"

# Create a collection in Weaviate which includes the required schema fields shown below.
echo '{
  "class": "'${WEAVIATE_COLLECTION_NAME}'",
  "properties": [
    { "name": "fileId", "dataType": [ "string" ] },
    { "name": "corpusId", "dataType": [ "string" ] },
    { "name": "chunkId", "dataType": [ "string" ] },
    { "name": "chunkDataType", "dataType": [ "string" ] },
    { "name": "chunkData", "dataType": [ "string" ] },
    { "name": "fileOriginalUri", "dataType": [ "string" ] }
  ]
}' | curl \
    -X POST \
    -H 'Content-Type: application/json' \
    -H "Authorization: Bearer "${WEAVIATE_API_KEY} \
    -d @- \
    ${HTTP_ENDPOINT_NAME}/v1/schema

Set up your Secret Manager

To set up your Secret Manager, you must enable Secret Manager, and set permissions.

Enable your Secret Manager

To enable your Secret Manager, do the following:

1. Go to the Secret Manager page.

   Go to Secret Manager

2. Click + Create Secret.
3. Enter the Name of your secret. Secret names can only contain English letters (A-Z), numbers (0-9), dashes (-), and underscores (_).
4. Specifying the following fields is optional:
   1. To upload the file with your secret, click Browse.
   2. Read the Replication policy.
   3. If you want to manually manage the locations for your secret, then check Manually manage locations for this secret. At least one region must be selected.
   4. Select your encryption option.
   5. If you want to manually set your rotation period, then check Set rotation period.
   6. If you want to specify Publish or subscribe topic(s) to receive event notifications, click Add topics.
   7. By default, the secret never expires. If you want to set an expiration date, then check Set expiration date.
   8. By default, secret versions are destroyed upon request. To delay the destruction of secret versions, check Set duration for delayed destruction.
   9. If you want to use labels to organize and categorize your secrets, then click + Add label.
   10. If you want to use annotations to attach non-identifying metadata to your secrets, then click + Add annotation.
5. Click Create secret.

Set permissions

1. You must grant Secret Manager permissions to your service account.

2. In the IAM & Admin section of your Google Cloud console, find your service account account, and click the pencil icon to edit.

3. In the Role field, select Secret Manager Secret Accessor.

This code sample demonstrates how to set up your Secret Manager.

# TODO(developer): Update the variables.
# Select a resource name for your Secret, which contains your API Key.
SECRET_NAME="MyWeaviateApiKeySecret"

# Create a secret in SecretManager.
curl "https://secretmanager.googleapis.com/v1/projects/${PROJECT_ID}/secrets?secretId=${SECRET_NAME}" \
    --request "POST" \
    --header "authorization: Bearer $(gcloud auth print-access-token)" \
    --header "content-type: application/json" \
    --data "{\"replication\": {\"automatic\": {}}}"

# Your Weaviate API Key.
WEAVIATE_API_KEY="example-api-key"
# Encode your WEAVIATE_API_KEY using base 64.
SECRET_DATA=$(echo ${WEAVIATE_API_KEY} | base64)

# Create a new version of your secret which uses SECRET_DATA as payload
curl.
"https://secretmanager.googleapis.com/v1/projects/${PROJECT_ID}/secrets/${SECRET_NAME}:addVersion" \
    --request "POST" \
    --header "authorization: Bearer $(gcloud auth print-access-token)" \
    --header "content-type: application/json" \
    --data "{\"payload\": {\"data\": \"${SECRET_DATA}\"}}"

Use Weaviate with Llama 3

The Model Garden RAG API notebook demonstrates how to use the Vertex AI SDK for Python with a Weaviate corpus and Llama 3 model. To use the notebook, you must do the following:

1. Set up your Weaviate database.

2. Set up your Secret Manager.

3. Use the Model Garden RAG API notebook.

For more examples, see Examples.

Create a RAG corpus

This code sample demonstrates how to create a RAG corpus, and sets the Weaviate instance as its vector database.

# TODO(developer): Update the variables.
PROJECT_ID = "YOUR_PROJECT_ID"
# The HTTPS/HTTP Weaviate endpoint you created during provisioning.
HTTP_ENDPOINT_NAME="https://your.weaviate.endpoint.com"

# Your Weaviate collection name, which roughly corresponds to a Vertex AI Knowledge Engine Corpus.
# For example, "MyCollectionName"
# Note that the first letter needs to be capitalized.
# Otherwise, Weaviate will capitalize it for you.
WEAVIATE_COLLECTION_NAME="MyCollectionName"

# The resource name of your Weaviate API Key your Secret.
SECRET_NAME="MyWeaviateApiKeySecret"
# The Secret Manager resource name containing the API Key for your Weaviate endpoint.
# For example, projects/{project}/secrets/{secret}/versions/latest
APIKEY_SECRET_VERSION="projects/${PROJECT_ID}/secrets/${SECRET_NAME}/versions/latest"

# Select a Corpus display name.
CORPUS_DISPLAY_NAME="SpecialCorpus"

# Call CreateRagCorpus API and set all Vector DB Config parameters for Weaviate to create a new corpus associated to your selected Weaviate collection.
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" : {
              "weaviate": {
                    "http_endpoint": '\""${HTTP_ENDPOINT_NAME}"\"',
                    "collection_name": '\""${WEAVIATE_COLLECTION_NAME}"\"'
              },
        "api_auth" : {
                "api_key_config": {
                      "api_key_secret_version": '\""${APIKEY_SECRET_VERSION}"\"'
                }
        }
      }
  }'

# TODO(developer): Update the variables.
# Get operation_id returned in CreateRagCorpus.
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}

# 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"

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="publishers/google/models/text-embedding-004"
)

# Configure a third-party model or a Google fine-tuned first-party model as a Vertex Endpoint resource
# See https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_e5.ipynb for
# deploying 3P embedding models to endpoints
ENDPOINT_ID = "your-model-endpoint-id"  # @param {type:"string"}
MODEL_ENDPOINT = "projects/{PROJECT_ID}/locations/us-central1/endpoints/{ENDPOINT_ID}"

embedding_model_config = rag.EmbeddingModelConfig(
    endpoint=MODEL_ENDPOINT,
)

# Configure a Weaviate Vector Database Instance for the corpus
WEAVIATE_HTTP_ENDPOINT = "weaviate-http-endpoint"  # @param {type:"string"}
COLLECTION_NAME = "weaviate-collection-name"  # @param {type:"string"}
API_KEY = "your-secret-manager-resource-name"  # @param {type:"string"}

vector_db = rag.Weaviate(
    weaviate_http_endpoint=WEAVIATE_HTTP_ENDPOINT,
    collection_name=COLLECTION_NAME,
    api_key=API_KEY,
)

# Name your corpus
DISPLAY_NAME = "your-corpus-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()

Use the RAG file

The RAG API handles the file upload, import, listing, and deletion.

POST https://LOCATION-aiplatform.googleapis.com/upload/v1beta1/projects/PROJECT_ID/locations/LOCATION/ragCorpora/RAG_CORPUS_ID/ragFiles:upload

{
 "rag_file": {
  "display_name": "FILE_DISPLAY_NAME",
  "description": "RAG_FILE_DESCRIPTION"
 }
}

curl -X POST \
     -H "Content-Type: application/json; charset=utf-8" \
     -d @INPUT_FILE \
     "https://LOCATION-aiplatform.googleapis.com/upload/v1beta1/projects/PROJECT_ID/locations/LOCATION/ragCorpora/RAG_CORPUS_ID/ragFiles:upload"

$headers = @{  }

Invoke-WebRequest `
    -Method POST `
    -Headers $headers `
    -ContentType: "application/json; charset=utf-8" `
    -InFile INPUT_FILE `
    -Uri "https://LOCATION-aiplatform.googleapis.com/upload/v1beta1/projects/PROJECT_ID/locations/LOCATION/ragCorpora/RAG_CORPUS_ID/ragFiles:upload" | Select-Object -Expand Content

from vertexai.preview import rag
import vertexai

# TODO(developer): Update and un-comment below lines
# PROJECT_ID = "your-project-id"
# corpus_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}"
# path = "path/to/local/file.txt"
# display_name = "file_display_name"
# description = "file description"

# Initialize Vertex AI API once per session
vertexai.init(project=PROJECT_ID, location="us-central1")

rag_file = rag.upload_file(
    corpus_name=corpus_name,
    path=path,
    display_name=display_name,
    description=description,
)
print(rag_file)
# RagFile(name='projects/[PROJECT_ID]/locations/us-central1/ragCorpora/1234567890/ragFiles/09876543',
#  display_name='file_display_name', description='file description')

Import RAG files

Files and folders can be imported from Drive or Cloud Storage.

POST https://LOCATION-aiplatform.googleapis.com/upload/v1beta1/projects/PROJECT_ID/locations/LOCATION/ragCorpora/RAG_CORPUS_ID/ragFiles:import

{
  "import_rag_files_config": {
    "gcs_source": {
      "uris": GCS_URIS
    },
    "google_drive_source": {
      "resource_ids": {
        "resource_id": DRIVE_RESOURCE_ID,
        "resource_type": DRIVE_RESOURCE_TYPE
      },
    }
  }
}

curl -X POST \
     -H "Content-Type: application/json; charset=utf-8" \
     -d @request.json \
     "https://LOCATION-aiplatform.googleapis.com/upload/v1beta1/projects/PROJECT_ID/locations/LOCATION/ragCorpora/RAG_CORPUS_ID/ragFiles:import"

$headers = @{  }

Invoke-WebRequest `
    -Method POST `
    -Headers $headers `
    -ContentType: "application/json; charset=utf-8" `
    -InFile request.json `
    -Uri "https://LOCATION-aiplatform.googleapis.com/upload/v1beta1/projects/PROJECT_ID/locations/LOCATION/ragCorpora/RAG_CORPUS_ID/ragFiles:import" | Select-Object -Expand Content

// Cloud Storage bucket/file location.
// Such as "gs://rag-e2e-test/"
GCS_URIS=YOUR_GCS_LOCATION

// Enter the QPM rate to limit RAG's access to your embedding model
// Example: 1000
EMBEDDING_MODEL_QPM_RATE=MAX_EMBEDDING_REQUESTS_PER_MIN_LIMIT

// ImportRagFiles
// Import a single Cloud Storage file or all files in a Cloud Storage bucket.
// Input: ENDPOINT, PROJECT_ID, RAG_CORPUS_ID, GCS_URIS
// Output: ImportRagFilesOperationMetadataNumber
// Use ListRagFiles to find the server-generated rag_file_id.
curl -X POST \
-H "Authorization: Bearer $(gcloud auth print-access-token)" \
-H "Content-Type: application/json" \
https://${ENDPOINT}/v1beta1/projects/${PROJECT_ID}/locations/${LOCATION}/ragCorpora/${RAG_CORPUS_ID}/ragFiles:import \
-d '{
  "import_rag_files_config": {
    "gcs_source": {
      "uris": '\""${GCS_URIS}"\"'
    },
    "rag_file_chunking_config": {
      "chunk_size": 512
    },
    "max_embedding_requests_per_min": '"${EMBEDDING_MODEL_QPM_RATE}"'
  }
}'

// Poll the operation status.
// The response contains the number of files imported.
OPERATION_ID=OPERATION_ID
poll_op_wait ${OPERATION_ID}

// Google Drive folder location.
FOLDER_RESOURCE_ID=YOUR_GOOGLE_DRIVE_FOLDER_RESOURCE_ID

// Enter the QPM rate to limit RAG's access to your embedding model
// Example: 1000
EMBEDDING_MODEL_QPM_RATE=MAX_EMBEDDING_REQUESTS_PER_MIN_LIMIT

// ImportRagFiles
// Import all files in a Google Drive folder.
// Input: ENDPOINT, PROJECT_ID, RAG_CORPUS_ID, FOLDER_RESOURCE_ID
// Output: ImportRagFilesOperationMetadataNumber
// Use ListRagFiles to find the server-generated rag_file_id.
curl -X POST \
-H "Authorization: Bearer $(gcloud auth print-access-token)" \
-H "Content-Type: application/json" \
https://${ENDPOINT}/v1beta1/projects/${PROJECT_ID}/locations/${LOCATION}/ragCorpora/${RAG_CORPUS_ID}/ragFiles:import \
-d '{
  "import_rag_files_config": {
    "google_drive_source": {
      "resource_ids": {
        "resource_id": '\""${FOLDER_RESOURCE_ID}"\"',
        "resource_type": "RESOURCE_TYPE_FOLDER"
      }
    },
    "max_embedding_requests_per_min": '"${EMBEDDING_MODEL_QPM_RATE}"'
  }
}'

// Poll the operation status.
// The response contains the number of files imported.
OPERATION_ID=OPERATION_ID
poll_op_wait ${OPERATION_ID}

from vertexai.preview import rag
import vertexai

# TODO(developer): Update and un-comment below lines
# PROJECT_ID = "your-project-id"
# corpus_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}"
# paths = ["https://drive.google.com/file/123", "gs://my_bucket/my_files_dir"]  # Supports Google Cloud Storage and Google Drive Links

# Initialize Vertex AI API once per session
vertexai.init(project=PROJECT_ID, location="us-central1")

response = rag.import_files(
    corpus_name=corpus_name,
    paths=paths,
    chunk_size=512,  # Optional
    chunk_overlap=100,  # Optional
    max_embedding_requests_per_min=900,  # Optional
)
print(f"Imported {response.imported_rag_files_count} files.")
# Example response:
# Imported 2 files.

Get a RAG file

GET https://LOCATION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/LOCATION/ragCorpora/RAG_CORPUS_ID/ragFiles/RAG_FILE_ID

curl -X GET \
     "https://LOCATION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/LOCATION/ragCorpora/RAG_CORPUS_ID/ragFiles/RAG_FILE_ID"

$headers = @{  }

Invoke-WebRequest `
    -Method GET `
    -Headers $headers `
    -Uri "https://LOCATION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/LOCATION/ragCorpora/RAG_CORPUS_ID/ragFiles/RAG_FILE_ID" | Select-Object -Expand Content

from vertexai.preview import rag
import vertexai

# TODO(developer): Update and un-comment below lines
# PROJECT_ID = "your-project-id"
# file_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}/ragFiles/{rag_file_id}"

# Initialize Vertex AI API once per session
vertexai.init(project=PROJECT_ID, location="us-central1")

rag_file = rag.get_file(name=file_name)
print(rag_file)
# Example response:
# RagFile(name='projects/1234567890/locations/us-central1/ragCorpora/11111111111/ragFiles/22222222222',
# display_name='file_display_name', description='file description')

List RAG files

GET https://LOCATION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/LOCATION/ragCorpora/RAG_CORPUS_ID/ragFiles?page_size=PAGE_SIZE&page_token=PAGE_TOKEN

curl -X GET \
     "https://LOCATION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/LOCATION/ragCorpora/RAG_CORPUS_ID/ragFiles?page_size=PAGE_SIZE&page_token=PAGE_TOKEN"

$headers = @{  }

Invoke-WebRequest `
    -Method GET `
    -Headers $headers `
    -Uri "https://LOCATION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/LOCATION/ragCorpora/RAG_CORPUS_ID/ragFiles?page_size=PAGE_SIZE&page_token=PAGE_TOKEN" | Select-Object -Expand Content

from vertexai.preview import rag
import vertexai

# TODO(developer): Update and un-comment below lines
# PROJECT_ID = "your-project-id"
# corpus_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}"

# Initialize Vertex AI API once per session
vertexai.init(project=PROJECT_ID, location="us-central1")

files = rag.list_files(corpus_name=corpus_name)
for file in files:
    print(file.display_name)
    print(file.name)
# Example response:
# g-drive_file.txt
# projects/1234567890/locations/us-central1/ragCorpora/111111111111/ragFiles/222222222222
# g_cloud_file.txt
# projects/1234567890/locations/us-central1/ragCorpora/111111111111/ragFiles/333333333333

Delete a RAG file

DELETE https://LOCATION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/LOCATION/ragCorpora/RAG_CORPUS_ID/ragFiles/RAG_FILE_ID

curl -X DELETE \
     "https://LOCATION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/LOCATION/ragCorpora/RAG_CORPUS_ID/ragFiles/RAG_FILE_ID"

$headers = @{  }

Invoke-WebRequest `
    -Method DELETE `
    -Headers $headers `
    -Uri "https://LOCATION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/LOCATION/ragCorpora/RAG_CORPUS_ID/ragFiles/RAG_FILE_ID" | Select-Object -Expand Content

from vertexai.preview import rag
import vertexai

# TODO(developer): Update and un-comment below lines
# PROJECT_ID = "your-project-id"
# file_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}/ragFiles/{rag_file_id}"

# Initialize Vertex AI API once per session
vertexai.init(project=PROJECT_ID, location="us-central1")

rag.delete_file(name=file_name)
print(f"File {file_name} deleted.")
# Example response:
# Successfully deleted the RagFile.
# File projects/1234567890/locations/us-central1/ragCorpora/1111111111/ragFiles/2222222222 deleted.

Retrieve context

When a user asks a question or provides a prompt, the retrieval component in RAG searches through its knowledge base to find information that is relevant to the query.

POST https://LOCATION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/LOCATION:retrieveContexts

{
 "vertex_rag_store": {
    "rag_resources": {
      "rag_corpus": "RAG_CORPUS_RESOURCE",
    },
    "vector_distance_threshold": 0.8
  },
  "query": {
   "text": "TEXT",
   "similarity_top_k": SIMILARITY_TOP_K
  }
 }

curl -X POST \
     -H "Content-Type: application/json; charset=utf-8" \
     -d @request.json \
     "https://LOCATION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/LOCATION:retrieveContexts"

$headers = @{  }

Invoke-WebRequest `
    -Method POST `
    -Headers $headers `
    -ContentType: "application/json; charset=utf-8" `
    -InFile request.json `
    -Uri "https://LOCATION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/LOCATION:retrieveContexts" | Select-Object -Expand Content

from vertexai.preview import rag
import vertexai

# TODO(developer): Update and un-comment below lines
# PROJECT_ID = "your-project-id"
# corpus_name = "projects/[PROJECT_ID]/locations/us-central1/ragCorpora/[rag_corpus_id]"

# Initialize Vertex AI API once per session
vertexai.init(project=PROJECT_ID, location="us-central1")

response = rag.retrieval_query(
    rag_resources=[
        rag.RagResource(
            rag_corpus=corpus_name,
            # Optional: supply IDs from `rag.list_files()`.
            # rag_file_ids=["rag-file-1", "rag-file-2", ...],
        )
    ],
    text="Hello World!",
    similarity_top_k=10,  # Optional
    vector_distance_threshold=0.5,  # Optional
)
print(response)
# Example response:
# contexts {
#   contexts {
#     source_uri: "gs://your-bucket-name/file.txt"
#     text: "....
#   ....

Generates content

A prediction controls the LLM method that generates content.

POST https://LOCATION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/LOCATION/publishers/google/models/MODEL_ID:GENERATION_METHOD

{
 "contents": {
  "role": "user",
  "parts": {
    "text": "INPUT_PROMPT"
  }
 },
 "tools": {
  "retrieval": {
   "disable_attribution": false,
   "vertex_rag_store": {
    "rag_resources": {
      "rag_corpus": "RAG_CORPUS_RESOURCE",
    },
    "similarity_top_k": SIMILARITY_TOP_K,
    "vector_distance_threshold": VECTOR_DISTANCE_THRESHOLD
   }
  }
 }
}

curl -X POST \
     -H "Content-Type: application/json; charset=utf-8" \
     -d @request.json \
     "https://LOCATION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/LOCATION/publishers/google/models/MODEL_ID:GENERATION_METHOD"

$headers = @{  }

Invoke-WebRequest `
    -Method POST `
    -Headers $headers `
    -ContentType: "application/json; charset=utf-8" `
    -InFile request.json `
    -Uri "https://LOCATION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/LOCATION/publishers/google/models/MODEL_ID:GENERATION_METHOD" | Select-Object -Expand Content

from vertexai.preview import rag
from vertexai.preview.generative_models import GenerativeModel, Tool
import vertexai

# TODO(developer): Update and un-comment below lines
# PROJECT_ID = "your-project-id"
# corpus_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}"

# Initialize Vertex AI API once per session
vertexai.init(project=PROJECT_ID, location="us-central1")

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

rag_model = GenerativeModel(
    model_name="gemini-1.5-flash-001", tools=[rag_retrieval_tool]
)
response = rag_model.generate_content("Why is the sky blue?")
print(response.text)
# Example response:
#   The sky appears blue due to a phenomenon called Rayleigh scattering.
#   Sunlight, which contains all colors of the rainbow, is scattered
#   by the tiny particles in the Earth's atmosphere....
#   ...

Hybrid search

Hybrid search is supported with Weaviate database, which combines both semantic and keyword searches to improve the relevance of search results. During the retrieval of search results, a combination of similarity scores from semantic (a dense vector) and keyword matching (a sparse vector) produces the final ranked results.

Hybrid search using the RAG Engine retrieval API

This is an example of how to enable a hybrid search using the RAG Engine retrieval API.

  POST https://us-central1-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/us-central1:retrieveContexts

  {
      "vertex_rag_store": {
        "rag_resources": {
            "rag_corpus": '\""${RAG_CORPUS_RESOURCE}"\"',
          },

        "vector_distance_threshold": ${DISTANCE_THRESHOLD}
      },
      "query": {
        "text": '\""${RETRIEVAL_QUERY}"\"',
        "similarity_top_k": ${TOP_K},
        "ranking": { "alpha" : ${ALPHA}}
      }
  }

from vertexai.preview import rag
import vertexai

# TODO(developer): Update the variables.
# PROJECT_ID = "your-project-id"
# rag_corpus_id = "your-rag-corpus-id"
# Only one corpus is supported at this time

# Initialize Vertex AI API once per session
vertexai.init(project=PROJECT_ID, location="us-central1")

response = rag.retrieval_query(
    rag_resources=[
        rag.RagResource(
            rag_corpus=rag_corpus_id,
            # Optional: supply IDs from `rag.list_files()`.
            # rag_file_ids=["rag-file-1", "rag-file-2", ...],
        )
    ],
    text="What is RAG and why it is helpful?",
    similarity_top_k=10,  # Optional
    vector_distance_threshold=0.5,  # Optional
    ranking=rag.RagQuery.Ranking(
                alpha=0.5
    ),  # Optional
)
print(response)

Use hybrid search and RAG Engine for grounded generation

This is an example of how to use hybrid search and RAG Engine for grounded generation.

    POST https://us-central1-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/us-central1/publishers/google/models/gemini-pro:generateContent

    {
      "contents": {
        "role": "user",
        "parts": {
          "text": '\""${INPUT_PROMPT}"\"'
        }
      },
      "tools": {
        "retrieval": {
          "vertex_rag_store": {
            "rag_resources": {
                "rag_corpus": '\""${RAG_CORPUS_RESOURCE}"\"',
              },
            "similarity_top_k": ${TOP_K},
            "vector_distance_threshold": ${DISTANCE_THRESHOLD},
            "ranking": { "alpha" : ${ALPHA}}
          }
        }
      }
    }

from vertexai.preview import rag
from vertexai.preview.generative_models import GenerativeModel, Tool
import vertexai

# TODO(developer): Update the variables.
# PROJECT_ID = "your-project-id"
# rag_corpus_id = "your-rag-corpus-id" # Only one corpus is supported at this time

# Initialize Vertex AI API once per session
vertexai.init(project=PROJECT_ID, location="us-central1")

rag_retrieval_tool = Tool.from_retrieval(
    retrieval=rag.Retrieval(
        source=rag.VertexRagStore(
            rag_resources=[
                rag.RagResource(
                    rag_corpus=rag_corpus_id,  # Currently only 1 corpus is allowed.
                    # Optional: supply IDs from `rag.list_files()`.
                    # rag_file_ids=["rag-file-1", "rag-file-2", ...],
                )
            ],
            similarity_top_k=3,  # Optional
            vector_distance_threshold=0.5,  # Optional
            ranking=rag.RagQuery.Ranking(
                        alpha=0.5
            ),  # Optional
        ),
    )
)

rag_model = GenerativeModel(
    model_name="gemini-1.5-flash-001", tools=[rag_retrieval_tool]
)
response = rag_model.generate_content("Why is the sky blue?")
print(response.text)

What's next

• To learn more about grounding, see Grounding overview.
• To learn more about RAG Engine, see Use RAG Engine.
• To learn more about grounding and RAG, see Ground responses using RAG.