rag-with-vectorsearch-ds

Agentic RAG Project with Vertex AI Vector Search and Firestore in Datastore mode

This project is a sample implementation of an Agentic RAG using the Agent Development Kit (ADK), with Vertex AI Vector Search as the vector store and Firestore in Datastore mode as the document store.

Project Structure

/rag-with-vectorsearch-ds
├── rag_with_vectorsearch_ds/    # ADK Agent directory
│   ├── .env.example
│   ├── agent.py
│   ├── tools.py
│   └── requirements.txt         # Agent dependencies
├── data_ingestion/              # Data ingestion directory
│   ├── .env.example
│   ├── ingest.py
│   ├── create_vector_search_index.py
│   └── requirements.txt         # Data ingestion script dependencies
├── source_documents/            # Source documents for RAG
└── README.md

Architecture Pattern: Using Vertex AI Vector Search (Specialized Service Architecture)

This pattern is used for handling extremely large-scale vector data (billions of vectors) or when ultra-low latency and advanced filtering capabilities are required. Vertex AI Vector Search is a dedicated, specialized service for large-scale similarity searches.

How It Works

• Store Ground Truth Data: The original data is still stored in Firestore to ensure data consistency and easy management.
• Generate Embeddings: Create vector embeddings of the data using the Vertex AI Embedding API.
• Build Index: Upload the generated vector embeddings and their corresponding data IDs to Vertex AI Vector Search to build a dedicated index.
• Primary Search (Get IDs): When a user makes a search request, the application first queries Vertex AI Vector Search to get a list of IDs for the most similar data points.
• Secondary Lookup (Get Data): The application then uses these returned IDs to look up the full, original data from Firestore to present to the user.

Architecture Diagram

+--------------+    (1) Query      +----------------------------+  (5) Fetch Data by ID +------------------------+
|              | ----------------> |        Agentic RAG         | --------------------> |  Firestore (Datastore) |
|  User/Client | <---------------- |(Cloud Run, Agent Engine...)| <-------------------- |   (Ground Truth Data)  |
|              | (6) Final Result  +----------------------------+  (5) Return Data      +------------------------+
+--------------+                          |            ^
                            (2) Embedding |            | (4) Return similar data IDs
                                          v            |
                             +------------------+  +---------------------------+
                             |  Vertex AI       |  | Vertex AI Vector Search   |
                             |  Embedding API   |  | (ScaNN-based Index)       |
                             +------------------+  +---------------------------+

Prerequisites

Before you begin, you need to have an active Google Cloud project.

1. Configure your Google Cloud project

First, you need to authenticate with Google Cloud. Run the following command and follow the instructions to log in.

gcloud auth application-default login

Next, set up your project and enable the necessary APIs.

# Set your project ID and location
export PROJECT_ID=$(gcloud config get-value project)
export LOCATION="us-central1" # Or your preferred location

# Enable the required APIs
gcloud services enable \
  aiplatform.googleapis.com \
  datastore.googleapis.com \
  cloudresourcemanager.googleapis.com

2. Create a Vector Search Index and Endpoint

This project provides a script to create the necessary Vector Search index and endpoint.

Run the index creation script: The create_vector_search_index.py script automates the creation of both the index and the endpoint. You must provide a name for your index. Note: This process can take 20-30 minutes to complete.

# Navigate to the data ingestion directory
cd rag-with-vectorsearch-ds/data_ingestion

# Run the script
python create_vector_search_index.py --index_name="my-rag-ds-index"

The script will output the Index ID and Endpoint ID. Make sure to save these values, as you will need them for the next steps.

3. Create a Firestore in Datastore mode database

Create a Firestore database in Datastore mode. This database will be used by the ingest.py script to store the document chunks.

export DATASTORE_DATABASE="vectorstore"
gcloud firestore databases create \
    --project=$PROJECT_ID \
    --location="nam5" \
    --database=$DATASTORE_DATABASE \
    --type=datastore-mode

4. Grant Agent Engine permissions

To allow the deployed Agent Engine to access your Vector Search index and Datastore, you must grant the Vertex AI User and Datastore User roles to the Agent Engine's service account.

Run the following commands to grant the necessary roles:

# Get your project number
export PROJECT_NUMBER=$(gcloud projects describe $PROJECT_ID --format="value(projectNumber)")

# Grant the Vertex AI User role
gcloud projects add-iam-policy-binding $PROJECT_ID \
    --member="serviceAccount:service-${PROJECT_NUMBER}@gcp-sa-aiplatform-re.iam.gserviceaccount.com" \
    --role="roles/aiplatform.user"

# Grant the Datastore User role
gcloud projects add-iam-policy-binding $PROJECT_ID \
    --member="serviceAccount:service-${PROJECT_NUMBER}@gcp-sa-aiplatform-re.iam.gserviceaccount.com" \
    --role="roles/datastore.user"

Without these permissions, you will encounter 403 IAM_PERMISSION_DENIED errors.

Setup

1. Install Dependencies

This project uses uv to manage the Python virtual environment and package dependencies.

Create and activate the virtual environment: From the root of the rag-with-vectorsearch-ds directory:

# Create the virtual environment
uv venv

# Activate the virtual environment (macOS/Linux)
source .venv/bin/activate
# Activate the virtual environment (Windows)
.venv\Scripts\activate

Install dependencies:

# Install agent dependencies
uv pip install -r rag_with_vectorsearch_ds/requirements.txt

# Install data ingestion script dependencies
uv pip install -r data_ingestion/requirements.txt

2. Data Ingestion

Run the data_ingestion/ingest.py script to load documents from source_documents/ into Datastore and create embeddings in Vertex AI Vector Search.

This script requires your GCP PROJECT_ID, LOCATION, INDEX_ID, and ENDPOINT_ID from the previous steps.

Example:

# Navigate to the data ingestion directory
cd rag-with-vectorsearch-ds/data_ingestion

# Run the data ingestion script
python ingest.py \
  --project_id="your-gcp-project-id" \
  --location="us-central1" \
  --index_id="your-vector-search-index-id" \
  --endpoint_id="your-vector-search-endpoint-id"

3. Run the Agent Locally

Before running the agent, you need to create a .env file in the rag_with_vectorsearch_ds directory. Copy the example file and fill in the required values for your environment.

cp rag_with_vectorsearch_ds/.env.example rag_with_vectorsearch_ds/.env
# Now, open rag_with_vectorsearch_ds/.env in an editor and modify the values.
# You will need to provide:
# - GOOGLE_CLOUD_PROJECT
# - GOOGLE_CLOUD_LOCATION
# - VECTOR_SEARCH_INDEX_ID
# - VECTOR_SEARCH_ENDPOINT_ID
# - DATASTORE_DATABASE (optional, defaults to "vectorstore")
# - DATASTORE_KIND (optional, defaults to "document_chunk")

You can run the agent using either the command-line interface or a web-based interface.

Using the Command-Line Interface (CLI)

Run the agent in your terminal using the adk run command.

adk run rag_with_vectorsearch_ds

Using the Web Interface

You can also interact with the agent through a web interface using the adk web command.

adk web

Screenshot:

Deployment

The RAG with Vector Search agent can be deployed to Vertex AI Agent Engine using the following commands.

1. Set Environment Variables

Before running the deployment script, you need to set the following environment variables.

export GOOGLE_CLOUD_PROJECT=$(gcloud config get-value project)
export GOOGLE_CLOUD_LOCATION="your-gcp-location"
export GOOGLE_CLOUD_STORAGE_BUCKET="your-gcs-bucket-for-staging"

2. Install Deployment Dependencies

You will need to install google-cloud-aiplatform with the agent_engines extra.

uv pip install "google-cloud-aiplatform[agent_engines]>=1.56.0" absl-py

3. Run the Deployment Script

Navigate to the deployment directory and run the script with the create action.

python deployment/deploy.py create

When the deployment finishes, it will print a line like this:

Successfully created remote agent: projects/<PROJECT_NUMBER>/locations/<PROJECT_LOCATION>/reasoningEngines/<AGENT_ENGINE_ID>

Make a note of the AGENT_ENGINE_ID. You will need it to interact with your deployed agent.

If you forgot the ID, you can list existing agents using:

python deployment/deploy.py list

To delete the deployed agent, you may run the following command:

python deployment/deploy.py delete --resource-id=${AGENT_ENGINE_ID}

4. Interact with the Deployed Agent

You can interact with your deployed agent using a simple Python script.

a. Set Environment Variables: Ensure the following environment variables are set in your terminal. You will need the AGENT_ENGINE_ID from the deployment step.

export GOOGLE_CLOUD_PROJECT="your-gcp-project-id"
export GOOGLE_CLOUD_LOCATION="your-gcp-location"
export AGENT_ENGINE_ID="your-agent-engine-id"

b. Create and Run the Python Script: Create a file named query_agent.py in the root of the project and add the following code.

import os
import vertexai
from vertexai import agent_engines

def query_remote_agent(project_id, location, agent_id, user_query):
    """Initializes Vertex AI and sends a query to the deployed agent."""
    vertexai.init(project=project_id, location=location)

    # Load the deployed agent
    remote_agent = agent_engines.get(agent_id)
    
    print(f"Querying agent: '{user_query}'...")

    # Stream the query and print the final text response
    response = remote_agent.query(
        message=user_query
    )
    print("Response:", response['output'])

if __name__ == "__main__":
    project = os.getenv("GOOGLE_CLOUD_PROJECT")
    loc = os.getenv("GOOGLE_CLOUD_LOCATION")
    agent = os.getenv("AGENT_ENGINE_ID")
    
    if not all([project, loc, agent]):
        print("Error: GOOGLE_CLOUD_PROJECT, GOOGLE_CLOUD_LOCATION, and AGENT_ENGINE_ID environment variables must be set.")
    else:
        query = "What is ADK?"
        query_remote_agent(project, loc, agent, query)

c. Run the script:

python query_agent.py

References

• Vertex AI Vector Search Quickstart
• Firestore in Datastore Mode Documentation
• Storing and querying data in Datastore
• Datastore Client Libraries
• LangChain VectorSearchVectorStore
• Google Cloud Generative AI Embeddings Samples