RAG quickstart for Python

This page shows you how to use the Vertex AI SDK to run Vertex AI RAG Engine tasks.

You can also follow along using this notebook Intro to Vertex AI RAG Engine .

Required roles

Grant roles to your user account. Run the following command once for each of the following IAM roles: roles/aiplatform.user

gcloud  
projects  
add-iam-policy-binding  
 PROJECT_ID 
  
--member = 
 "user: USER_IDENTIFIER 
" 
  
--role = 
 ROLE 

Replace the following:

  • PROJECT_ID : your project ID.
  • USER_IDENTIFIER : the identifier for your user account—for example, myemail@example.com .
  • ROLE : the IAM role that you grant to your user account.

Prepare your Google Cloud console

To use Vertex AI RAG Engine, do the following:

  1. Install the Vertex AI SDK for Python .

  2. Run this command in the Google Cloud console to set up your project.

    gcloud config set {project}

  3. Run this command to authorize your login.

    gcloud auth application-default login

Run Vertex AI RAG Engine

Copy and paste this sample code into the Google Cloud console to run Vertex AI RAG Engine.

Python

To learn how to install or update the Vertex AI SDK for Python, see Install the Vertex AI SDK for Python . For more information, see the Python API reference documentation .

  from 
  
 vertexai 
  
 import 
 rag 
 from 
  
 vertexai.generative_models 
  
 import 
 GenerativeModel 
 , 
 Tool 
 import 
  
 vertexai 
 # Create a RAG Corpus, Import Files, and Generate a response 
 # TODO(developer): Update and un-comment below lines 
 # PROJECT_ID = "your-project-id" 
 # display_name = "test_corpus" 
 # paths = ["https://drive.google.com/file/d/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" 
 ) 
 # Create RagCorpus 
 # Configure embedding model, for example "text-embedding-005". 
 embedding_model_config 
 = 
 rag 
 . 
 RagEmbeddingModelConfig 
 ( 
 vertex_prediction_endpoint 
 = 
 rag 
 . 
 VertexPredictionEndpoint 
 ( 
 publisher_model 
 = 
 "publishers/google/models/text-embedding-005" 
 ) 
 ) 
 rag_corpus 
 = 
 rag 
 . 
 create_corpus 
 ( 
 display_name 
 = 
 display_name 
 , 
 backend_config 
 = 
 rag 
 . 
 RagVectorDbConfig 
 ( 
 rag_embedding_model_config 
 = 
 embedding_model_config 
 ), 
 ) 
 # Import Files to the RagCorpus 
 rag 
 . 
 import_files 
 ( 
 rag_corpus 
 . 
 name 
 , 
 paths 
 , 
 # Optional 
 transformation_config 
 = 
 rag 
 . 
 TransformationConfig 
 ( 
 chunking_config 
 = 
 rag 
 . 
 ChunkingConfig 
 ( 
 chunk_size 
 = 
 512 
 , 
 chunk_overlap 
 = 
 100 
 , 
 ), 
 ), 
 max_embedding_requests_per_min 
 = 
 1000 
 , 
 # Optional 
 ) 
 # Direct context retrieval 
 rag_retrieval_config 
 = 
 rag 
 . 
 RagRetrievalConfig 
 ( 
 top_k 
 = 
 3 
 , 
 # Optional 
 filter 
 = 
 rag 
 . 
 Filter 
 ( 
 vector_distance_threshold 
 = 
 0.5 
 ), 
 # Optional 
 ) 
 response 
 = 
 rag 
 . 
 retrieval_query 
 ( 
 rag_resources 
 = 
 [ 
 rag 
 . 
 RagResource 
 ( 
 rag_corpus 
 = 
 rag_corpus 
 . 
 name 
 , 
 # 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?" 
 , 
 rag_retrieval_config 
 = 
 rag_retrieval_config 
 , 
 ) 
 print 
 ( 
 response 
 ) 
 # Enhance generation 
 # Create a RAG retrieval tool 
 rag_retrieval_tool 
 = 
 Tool 
 . 
 from_retrieval 
 ( 
 retrieval 
 = 
 rag 
 . 
 Retrieval 
 ( 
 source 
 = 
 rag 
 . 
 VertexRagStore 
 ( 
 rag_resources 
 = 
 [ 
 rag 
 . 
 RagResource 
 ( 
 rag_corpus 
 = 
 rag_corpus 
 . 
 name 
 , 
 # Currently only 1 corpus is allowed. 
 # Optional: supply IDs from `rag.list_files()`. 
 # rag_file_ids=["rag-file-1", "rag-file-2", ...], 
 ) 
 ], 
 rag_retrieval_config 
 = 
 rag_retrieval_config 
 , 
 ), 
 ) 
 ) 
 # Create a Gemini model instance 
 rag_model 
 = 
 GenerativeModel 
 ( 
 model_name 
 = 
 "gemini-2.0-flash-001" 
 , 
 tools 
 = 
 [ 
 rag_retrieval_tool 
 ] 
 ) 
 # Generate response 
 response 
 = 
 rag_model 
 . 
 generate_content 
 ( 
 "What is RAG and why it is helpful?" 
 ) 
 print 
 ( 
 response 
 . 
 text 
 ) 
 # Example response: 
 #   RAG stands for Retrieval-Augmented Generation. 
 #   It's a technique used in AI to enhance the quality of responses 
 # ... 
 

What's next

Create a Mobile Website
View Site in Mobile | Classic
Share by: