vertex-ai-rag-extension
Vertex AI RAG Engine Extension for Gemini CLI
Open source Open in the app JSON README (API)
About
Vertex AI RAG Engine Extension for Gemini CLI
Details
- Kind
- Plugins
- Topic
- AI, RAG & memory
- Publisher
- skpathak2
- Origin
- gemini
- Category
- ferramentas
- Version
- 1.0.0
- Stars
- 4
- Last push
- 2026-03-20T13:53:38Z
- Repository state
- ativo
- Language
- TypeScript
- License
- Apache-2.0
- Added
- 2026-08-30 14:13:39
- Updated
- 2026-08-30 14:13:39
- Origin id
skpathak2/vertex-ai-rag-engine-extension
README
# Vertex AI RAG Engine Extension for Gemini CLI
Credits:- Mandeep Singh Bawa, Shlok Karpathak
## Overview
The **Vertex AI RAG Engine Extension** allows the Gemini CLI to interact with the Google Cloud Vertex AI Retrieval-Augmented Generation (RAG) Engine. With this extension, you can manage document corpus, import files, and retrieve context for answering questions based on your uploaded documents.
## Prerequisites
- **Node.js**: v18+
- **Gemini CLI**: Installed and configured.
- **Google Cloud Project**: A GCP project with Vertex AI API enabled.
- **Authentication**: Local credentials set up (e.g., Application Default Credentials).
## Installation
To install or link the extension for local development:
1. Navigate to the extension directory.
2. Build the project:
```bash
npm run build
```
3. Link the extension to the Gemini CLI:
```bash
gemini extensions link .
```
*During installation, you will be prompted to configure settings like Project ID.*
*Note: If you get an error that the extension is already installed, run `gemini extensions uninstall vertex-ai-rag-extension` first.*
## Configuration
The extension supports the following settings, which can be configured during installation or via the CLI:
| Setting Name | Environment Variable | Description | Default |
| :--- | :--- | :--- | :--- |
| **Project ID** | `PROJECT_ID` | Google Cloud Project ID for Vertex AI requests. | *Required* |
| **Location** | `LOCATION` | Google Cloud location/region (e.g., `us-central1`). | `us-central1` |
| **Enable Disruptive Actions** | `ENABLE_DISRUPTIVE_ACTIONS` | Set to `true` to enable `delete` tools. | `false` |
To change configuration:
```bash
gemini extensions config vertex-ai-rag-extension
```
## Capabilities & Tools
### Corpus Management
* **`list_rag_corpus`**: Lists all RAG Corpus.
* *Params*: `page_size` (max items), `page_token` (pagination token).
* **`get_rag_corpus`**: Gets details of a specific corpus.
* *Params*: `rag_corpus_id` (required).
* **`create_rag_corpus`**: Creates a new RAG Corpus.
* *Params*: `display_name` (required), `description`.
* **`delete_rag_corpus`**: Deletes a corpus.
* *Safety*: Requires `Enable Disruptive Actions` to be true.
* *Params*: `rag_corpus_id` (required).
### File Management
* **`list_rag_files`**: Lists files in a corpus.
* *Params*: `rag_corpus_id` (required).
* **`get_rag_file`**: Gets file details.
* *Params*: `rag_corpus_id` (required), `rag_file_id` (required).
* **`import_rag_files`**: Imports files from GCS or Google Drive (Returns Operation).
* *Params*: `rag_corpus_id` (required), `import_config` (GCS or Drive Source).
* **`delete_rag_file`**: Deletes a file.
* *Safety*: Requires `Enable Disruptive Actions` to be true.
* *Params*: `rag_corpus_id` (required), `rag_file_id` (required).
### Retrieval (Core)
* **`retrieve_contexts`**: Retrieves relevant text snippets based on a query.
* *Params*:
* `rag_corpus_id` (required): Target corpus.
* `query_text` (required): Query text.
* `similarity_top_k`: Max snippets retrieved.
* `vector_distance_threshold`: Similarity threshold.
## Testing
An isolated integration test script is included to test the MCP server capabilities directly against the Vertex AI API without the CLI wrapper.
To run the integration tests:
1. Ensure you have authenticated your Google Cloud CLI:
```bash
gcloud auth application-default login
```
2. Run the test script:
```bash
npm run test:integration
```
This will run through a complete testing lifecycle including listing, creating, getting, and attempting to delete a RAG corpus to verify permissions, configuration, and API connectivity.
## Usage with Gemini CLI
Once installed and configured, the tools are available to the Gemini CLI models. You can ask the CLI to use the extension directly in prompts:
```bash
# Example: Using retrieval tool for context
gemini "Use vertex-ai-rag-extension to list all corpus"
# Example: Ask questions using an existing Corpus ID
gemini "Ask my corpus '123456789' what the safety policy is"
```
The model will automatically invoke the relevant sub-tools (like `retrieve_contexts`) as needed to fulfill the request.