From f791c2cfc567e82bbf95c216952b93755ecef523 Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Wed, 1 Jul 2026 11:19:54 -0400 Subject: [PATCH 01/15] docs: update RAG quickstart python snippet --- generative_ai/rag/quickstart_example.py | 122 ++++++++++++++---------- 1 file changed, 70 insertions(+), 52 deletions(-) diff --git a/generative_ai/rag/quickstart_example.py b/generative_ai/rag/quickstart_example.py index 32649f64aeb..1feecc6987d 100644 --- a/generative_ai/rag/quickstart_example.py +++ b/generative_ai/rag/quickstart_example.py @@ -13,11 +13,13 @@ # limitations under the License. import os +import agentplatform from typing import List, Tuple -from vertexai import rag -from vertexai.generative_models import GenerationResponse +from agentplatform import types +from google import genai +from google.genai import types as genai_types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") @@ -27,75 +29,85 @@ def quickstart( paths: List[str], ) -> Tuple[rag.RagCorpus, GenerationResponse]: # [START generativeaionvertexai_rag_quickstart] - from vertexai import rag - from vertexai.generative_models import GenerativeModel, Tool - import vertexai + import agentplatform + + from agentplatform import types + from google import genai + from google.genai import types as genai_types + # Create a RAG Corpus, Import Files, and Generate a response # TODO(developer): Update and un-comment below lines # PROJECT_ID = "your-project-id" + # MODEL_ID = "gemini-3.5-flash" # 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 + # gcs_path = "gs://my_bucket/my_files_dir/*" + # google_drive_path ="https://drive.google.com/file/d/123" # Initialize Vertex AI API once per session - vertexai.init(project=PROJECT_ID, location="us-east4") + client = agentplatform.Client(project=PROJECT_ID, location="us-east4") # 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 = client.rag.create_corpus( + rag_corpus=types.RagCorpus( + display_name=display_name, ) ) - 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 + client.rag.import_files( + name=rag_corpus.name, + import_config=types.ImportRagFilesConfig( + gcs_source=genai_types.GcsSource(uris=[gcs_path]), + google_drive_source=types.GoogleDriveSource( + resource_ids=[ + types.GoogleDriveSourceResourceId( + resource_id=google_drive_path, + resource_type=types.ResourceType.RESOURCE_TYPE_FILE + ) + ] + ), # optional + rag_file_transformation_config=types.RagFileTransformationConfig( + rag_file_chunking_config=types.RagFileChunkingConfig( + chunk_size=512, + chunk_overlap=100, + ) + ), # optional + max_embedding_requests_per_min=1000, # optional + ) ) # Direct context retrieval - rag_retrieval_config = rag.RagRetrievalConfig( + rag_retrieval_config = genai_types.RagRetrievalConfig( top_k=3, # Optional - filter=rag.Filter(vector_distance_threshold=0.5), # Optional + filter=genai_types.RagRetrievalConfigFilter( + 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, + response = client.rag.retrieve_contexts( + vertex_rag_store=genai_types.VertexRagStore( + rag_resources=[ + genai_types.VertexRagStoreRagResource( + rag_corpus=rag_corpus.name, + ) + ], + ), + query=types.RagQuery( + 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_retrieval_tool = genai_types.Tool( + retrieval=genai_types.Retrieval( + source=genai_types.VertexRagStore( rag_resources=[ - rag.RagResource( - rag_corpus=rag_corpus.name, # Currently only 1 corpus is allowed. + genai_types.VertexRagStoreRagResource( + rag_corpus=rag_corpus.name, # Optional: supply IDs from `rag.list_files()`. # rag_file_ids=["rag-file-1", "rag-file-2", ...], ) @@ -105,13 +117,19 @@ def quickstart( ) ) - # Create a Gemini model instance - rag_model = GenerativeModel( - model_name="gemini-2.0-flash-001", tools=[rag_retrieval_tool] - ) + # Call generate_content with the tool using the GenAI SDK - # Generate response - response = rag_model.generate_content("What is RAG and why it is helpful?") + # Create a GenAI SDK client + genai_client = genai.Client(enterprise=True, project=PROJECT_ID, location="us-east4") + + + response = genai_client.models.generate_content( + model=MODEL_ID, + contents="What is RAG and why it is helpful?", + config=genai_types.GenerateContentConfig( + tools=[rag_retrieval_tool] + ) + ) print(response.text) # Example response: # RAG stands for Retrieval-Augmented Generation. From 4af7bc90359f4096d0c5831ed68fd92fa2064d03 Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Wed, 1 Jul 2026 11:28:37 -0400 Subject: [PATCH 02/15] update arg name --- generative_ai/rag/quickstart_example.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/generative_ai/rag/quickstart_example.py b/generative_ai/rag/quickstart_example.py index 1feecc6987d..2d8eddee127 100644 --- a/generative_ai/rag/quickstart_example.py +++ b/generative_ai/rag/quickstart_example.py @@ -104,7 +104,7 @@ def quickstart( # Create a RAG retrieval tool rag_retrieval_tool = genai_types.Tool( retrieval=genai_types.Retrieval( - source=genai_types.VertexRagStore( + vertex_rag_store=genai_types.VertexRagStore( rag_resources=[ genai_types.VertexRagStoreRagResource( rag_corpus=rag_corpus.name, From 4b418088c381967dd9d9843e567dcb396749c9a0 Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Wed, 1 Jul 2026 12:32:20 -0400 Subject: [PATCH 03/15] update additional rag samples --- generative_ai/rag/create_corpus_example.py | 32 +++++++------ .../create_corpus_feature_store_example.py | 35 ++++++++------ generative_ai/rag/delete_corpus_example.py | 9 ++-- generative_ai/rag/delete_file_example.py | 9 ++-- generative_ai/rag/get_corpus_example.py | 9 ++-- generative_ai/rag/get_file_example.py | 9 ++-- .../rag/import_files_async_example.py | 46 +++++++++++------- generative_ai/rag/import_files_example.py | 47 +++++++++++++------ generative_ai/rag/list_corpora_example.py | 9 ++-- generative_ai/rag/list_files_example.py | 11 ++--- 10 files changed, 126 insertions(+), 90 deletions(-) diff --git a/generative_ai/rag/create_corpus_example.py b/generative_ai/rag/create_corpus_example.py index 90b1aa60401..c68801bc347 100644 --- a/generative_ai/rag/create_corpus_example.py +++ b/generative_ai/rag/create_corpus_example.py @@ -15,7 +15,6 @@ from typing import Optional -from vertexai.preview.rag import RagCorpus PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") @@ -26,30 +25,35 @@ def create_corpus( ) -> RagCorpus: # [START generativeaionvertexai_rag_create_corpus] - from vertexai import rag - import vertexai + import agentplatform + from agentplatform import types # TODO(developer): Update and un-comment below lines # PROJECT_ID = "your-project-id" # display_name = "test_corpus" # description = "Corpus Description" - # Initialize Vertex AI API once per session - vertexai.init(project=PROJECT_ID, location="us-central1") + # Initialize Agent Platform client once per session + client = agentplatform.Client(project=PROJECT_ID, location="us-central1") - # Configure backend_config - backend_config = rag.RagVectorDbConfig( - rag_embedding_model_config=rag.RagEmbeddingModelConfig( - vertex_prediction_endpoint=rag.VertexPredictionEndpoint( - publisher_model="publishers/google/models/text-embedding-005" + # Configure project-level config + backend_config = types.RagVectorDbConfig( + rag_embedding_model_config=types.RagEmbeddingModelConfig( + vertex_prediction_endpoint=types.RagEmbeddingModelConfigVertexPredictionEndpoint( + endpoint="publishers/google/models/text-embedding-005" ) ) ) + client.rag.update_config( + updated_config=backend_config + ) - corpus = rag.create_corpus( - display_name=display_name, - description=description, - backend_config=backend_config, + # Create a corpus + corpus = client.rag.create_corpus( + rag_corpus=types.RagCorpus( + display_name=display_name, + description=description, + ) ) print(corpus) # Example response: diff --git a/generative_ai/rag/create_corpus_feature_store_example.py b/generative_ai/rag/create_corpus_feature_store_example.py index 8674887c1fe..b0422f009d2 100644 --- a/generative_ai/rag/create_corpus_feature_store_example.py +++ b/generative_ai/rag/create_corpus_feature_store_example.py @@ -15,8 +15,6 @@ from typing import Optional -from vertexai.preview.rag import RagCorpus - PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") @@ -27,8 +25,8 @@ def create_corpus_feature_store( ) -> RagCorpus: # [START generativeaionvertexai_rag_create_corpus_feature_store] - from vertexai.preview import rag - import vertexai + import agentplatform + from agentplatform import types # TODO(developer): Update and un-comment below lines # PROJECT_ID = "your-project-id" @@ -36,22 +34,31 @@ def create_corpus_feature_store( # display_name = "test_corpus" # description = "Corpus Description" - # Initialize Vertex AI API once per session - vertexai.init(project=PROJECT_ID, location="us-central1") + # Initialize Agent Platform client once per session + client = agentplatform.Client(project=PROJECT_ID, location="us-central1") # Configure embedding model (Optional) - embedding_model_config = rag.EmbeddingModelConfig( - publisher_model="publishers/google/models/text-embedding-004" + backend_config = types.RagVectorDbConfig( + rag_embedding_model_config=types.RagEmbeddingModelConfig( + vertex_prediction_endpoint=types.RagEmbeddingModelConfigVertexPredictionEndpoint( + endpoint="publishers/google/models/text-embedding-005" + ), + ), + vertex_feature_store=types.RagDbConfigVertexFeatureStore( + feature_view_resource_name=feature_view_name + ) ) # Configure Vector DB - vector_db = rag.VertexFeatureStore(resource_name=feature_view_name) + client.rag.update_config( + updated_config=backend_config + ) - corpus = rag.create_corpus( - display_name=display_name, - description=description, - embedding_model_config=embedding_model_config, - vector_db=vector_db, + corpus = client.rag.create_corpus( + rag_corpus=types.RagCorpus( + display_name=display_name, + description=description, + ) ) print(corpus) # Example response: diff --git a/generative_ai/rag/delete_corpus_example.py b/generative_ai/rag/delete_corpus_example.py index 4255110fe14..527126a92a7 100644 --- a/generative_ai/rag/delete_corpus_example.py +++ b/generative_ai/rag/delete_corpus_example.py @@ -20,17 +20,16 @@ def delete_corpus(corpus_name: str) -> None: # [START generativeaionvertexai_rag_delete_corpus] - from vertexai import rag - import vertexai + import agentplatform # 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") + # Initialize Agent Platform client once per session + client = agentplatform.Client(project=PROJECT_ID, location="us-central1") - rag.delete_corpus(name=corpus_name) + client.rag.delete_corpus(name=corpus_name) print(f"Corpus {corpus_name} deleted.") # Example response: # Successfully deleted the RagCorpus. diff --git a/generative_ai/rag/delete_file_example.py b/generative_ai/rag/delete_file_example.py index e11afc71d96..477f93a97d0 100644 --- a/generative_ai/rag/delete_file_example.py +++ b/generative_ai/rag/delete_file_example.py @@ -20,17 +20,16 @@ def delete_file(file_name: str) -> None: # [START generativeaionvertexai_rag_delete_file] - from vertexai import rag - import vertexai + import agentplatform # 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") + # Initialize Agent Platform client once per session + client = agentplatform.Client(project=PROJECT_ID, location="us-central1") - rag.delete_file(name=file_name) + client.rag.delete_file(name=file_name) print(f"File {file_name} deleted.") # Example response: # Successfully deleted the RagFile. diff --git a/generative_ai/rag/get_corpus_example.py b/generative_ai/rag/get_corpus_example.py index 849995156d0..17037287d5c 100644 --- a/generative_ai/rag/get_corpus_example.py +++ b/generative_ai/rag/get_corpus_example.py @@ -22,17 +22,16 @@ def get_corpus(corpus_name: str) -> RagCorpus: # [START generativeaionvertexai_rag_get_corpus] - from vertexai import rag - import vertexai + import agentplatform # 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") + # Initialize Agent Platform client once per session + client = agentplatform.Client(project=PROJECT_ID, location="us-central1") - corpus = rag.get_corpus(name=corpus_name) + corpus = client.rag.get_corpus(name=corpus_name) print(corpus) # Example response: # RagCorpus(name='projects/[PROJECT_ID]/locations/us-central1/ragCorpora/1234567890', diff --git a/generative_ai/rag/get_file_example.py b/generative_ai/rag/get_file_example.py index 90c461ae4d9..afdfa7ce625 100644 --- a/generative_ai/rag/get_file_example.py +++ b/generative_ai/rag/get_file_example.py @@ -22,17 +22,16 @@ def get_file(file_name: str) -> RagFile: # [START generativeaionvertexai_rag_get_file] - from vertexai import rag - import vertexai + import agentplatform # 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") + # Initialize Agent Platform client once per session + client = agentplatform.Client(project=PROJECT_ID, location="us-central1") - rag_file = rag.get_file(name=file_name) + rag_file = client.rag.get_file(name=file_name) print(rag_file) # Example response: # RagFile(name='projects/1234567890/locations/us-central1/ragCorpora/11111111111/ragFiles/22222222222', diff --git a/generative_ai/rag/import_files_async_example.py b/generative_ai/rag/import_files_async_example.py index 7485b951ff0..9b00e0482cb 100644 --- a/generative_ai/rag/import_files_async_example.py +++ b/generative_ai/rag/import_files_async_example.py @@ -27,30 +27,44 @@ async def import_files_async( ) -> ImportRagFilesResponse: # [START generativeaionvertexai_rag_import_files_async] - from vertexai import rag - import vertexai + import agentplatform + from agentplatform import types + + from google.genai import types as genai_types # 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}" # Supports Google Cloud Storage and Google Drive Links - # paths = ["https://drive.google.com/file/d/123", "gs://my_bucket/my_files_dir"] - - # Initialize Vertex AI API once per session - vertexai.init(project=PROJECT_ID, location="us-central1") - - response = await rag.import_files( - corpus_name=corpus_name, - paths=paths, - transformation_config=rag.TransformationConfig( - rag.ChunkingConfig(chunk_size=512, chunk_overlap=100) - ), - max_embedding_requests_per_min=900, # Optional + # paths = ["https://drive.google.com/file/d/123", "gs://my_bucket/my_files_dir/*"] + + # Initialize Agent Platform client once per session + client = agentplatform.Client(project=PROJECT_ID, location="us-central1") + + response = await client.aio.rag.import_files( + name=corpus_name, + import_config=types.ImportRagFilesConfig( + gcs_source=genai_types.GcsSource(uris=[paths[1]]), + google_drive_source=types.GoogleDriveSource( + resource_ids=[ + types.GoogleDriveSourceResourceId( + resource_id=paths[0], + resource_type=types.ResourceType.RESOURCE_TYPE_FILE + ) + ] + ), # optional + rag_file_transformation_config=types.RagFileTransformationConfig( + rag_file_chunking_config=types.RagFileChunkingConfig( + chunk_size=512, + chunk_overlap=100, + ) + ), # optional + max_embedding_requests_per_min=900, # optional + ) ) - result = await response.result() - print(f"Imported {result.imported_rag_files_count} files.") + print(f"Imported {response.imported_rag_files_count} files.") # Example response: # Imported 2 files. diff --git a/generative_ai/rag/import_files_example.py b/generative_ai/rag/import_files_example.py index c21f68c28d2..36bbc08af9b 100644 --- a/generative_ai/rag/import_files_example.py +++ b/generative_ai/rag/import_files_example.py @@ -26,26 +26,43 @@ def import_files( ) -> ImportRagFilesResponse: # [START generativeaionvertexai_rag_import_files] - from vertexai import rag - import vertexai + import agentplatform + from agentplatform import types + + from google.genai import types as genai_types # 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, - transformation_config=rag.TransformationConfig( - rag.ChunkingConfig(chunk_size=512, chunk_overlap=100) - ), - import_result_sink="gs://sample-existing-folder/sample_import_result_unique.ndjson", # Optional, this has to be an existing storage bucket folder, and file name has to be unique (non-existent). - max_embedding_requests_per_min=900, # Optional + + # Supports Google Cloud Storage and Google Drive Links + # paths = ["https://drive.google.com/file/d/123", "gs://my_bucket/my_files_dir/*"] + + # Initialize Agent Platform client once per session + client = agentplatform.Client(project=PROJECT_ID, location="us-central1") + + response = client.rag.import_files( + name=corpus_name, + import_config=types.ImportRagFilesConfig( + gcs_source=genai_types.GcsSource(uris=[paths[1]]), + google_drive_source=types.GoogleDriveSource( + resource_ids=[ + types.GoogleDriveSourceResourceId( + resource_id=paths[0], + resource_type=types.ResourceType.RESOURCE_TYPE_FILE + ) + ] + ), # optional + rag_file_transformation_config=types.RagFileTransformationConfig( + rag_file_chunking_config=types.RagFileChunkingConfig( + chunk_size=512, + 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. diff --git a/generative_ai/rag/list_corpora_example.py b/generative_ai/rag/list_corpora_example.py index 138a47f4330..9bdab807b5f 100644 --- a/generative_ai/rag/list_corpora_example.py +++ b/generative_ai/rag/list_corpora_example.py @@ -24,16 +24,15 @@ def list_corpora() -> ListRagCorporaPager: # [START generativeaionvertexai_rag_list_corpora] - from vertexai import rag - import vertexai + import agentplatform # TODO(developer): Update and un-comment below lines # PROJECT_ID = "your-project-id" - # Initialize Vertex AI API once per session - vertexai.init(project=PROJECT_ID, location="us-central1") + # Initialize Agent Platform client once per session + client = agentplatform.Client(project=PROJECT_ID, location="us-central1") - corpora = rag.list_corpora() + corpora = client.rag.list_corpora() print(corpora) # Example response: # ListRagCorporaPager ListRagFilesPager: # [START generativeaionvertexai_rag_list_files] - from vertexai import rag - import vertexai + import agentplatform # 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") + # Initialize Agent Platform client once per session + client = agentplatform.Client(project=PROJECT_ID, location="us-central1") - files = rag.list_files(corpus_name=corpus_name) - for file in files: + files_response = client.rag.list_files(name=corpus_name) + for file in files_response.rag_files: print(file.display_name) print(file.name) # Example response: From bb9cbc7d103a8f244ff54764d227b5ecb9a0f4ef Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Wed, 1 Jul 2026 13:06:57 -0400 Subject: [PATCH 04/15] update remaining samples to use latest rag module --- generative_ai/rag/generate_content_example.py | 43 +++++++++++-------- generative_ai/rag/quickstart_example.py | 2 +- generative_ai/rag/retrieval_query_example.py | 41 +++++++++++------- generative_ai/rag/upload_file_example.py | 13 +++--- 4 files changed, 56 insertions(+), 43 deletions(-) diff --git a/generative_ai/rag/generate_content_example.py b/generative_ai/rag/generate_content_example.py index a02b8bfb7f1..4fde2cb2256 100644 --- a/generative_ai/rag/generate_content_example.py +++ b/generative_ai/rag/generate_content_example.py @@ -14,8 +14,6 @@ import os -from vertexai.generative_models import GenerationResponse - PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") @@ -24,39 +22,48 @@ def generate_content_with_rag( ) -> GenerationResponse: # [START generativeaionvertexai_rag_generate_content] - from vertexai import rag - from vertexai.generative_models import GenerativeModel, Tool - import vertexai + import agentplatform + + from agentplatform import types + from google import genai + from google.genai import types as genai_types # 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") + # Initialize Agent Platform client once per session + client = agentplatform.Client(project=PROJECT_ID, location="us-central1") - rag_retrieval_tool = Tool.from_retrieval( - retrieval=rag.Retrieval( - source=rag.VertexRagStore( + rag_retrieval_tool = genai_types.Tool( + retrieval=genai_types.Retrieval( + vertex_rag_store=genai_types.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", ...], - ) + genai_types.VertexRagStoreRagResource ], - rag_retrieval_config=rag.RagRetrievalConfig( + rag_retrieval_config=genai_types.RagRetrievalConfig( top_k=10, - filter=rag.utils.resources.Filter(vector_distance_threshold=0.5), + filter=genai_types.RagRetrievalConfigFilter( + vector_distance_threshold=0.5 + ), ), ), ) ) + # Create a GenAI SDK client to make a generate_content request + genai_client = genai.Client(enterprise=True, project=PROJECT_ID, location="us-central1") + rag_model = GenerativeModel( model_name="gemini-2.0-flash-001", tools=[rag_retrieval_tool] ) - response = rag_model.generate_content("Why is the sky blue?") + response = client.models.generate_content( + model="gemini-2.5-pro", + contents="Why is the sky blue?", + config=genai_types.GenerateContentConfig( + tools=[rag_retrieval_tool] + ) + ) print(response.text) # Example response: # The sky appears blue due to a phenomenon called Rayleigh scattering. diff --git a/generative_ai/rag/quickstart_example.py b/generative_ai/rag/quickstart_example.py index 2d8eddee127..f8fcf6af85c 100644 --- a/generative_ai/rag/quickstart_example.py +++ b/generative_ai/rag/quickstart_example.py @@ -45,7 +45,7 @@ def quickstart( # gcs_path = "gs://my_bucket/my_files_dir/*" # google_drive_path ="https://drive.google.com/file/d/123" - # Initialize Vertex AI API once per session + # Initialize Agent Platform client once per session client = agentplatform.Client(project=PROJECT_ID, location="us-east4") # Create RagCorpus diff --git a/generative_ai/rag/retrieval_query_example.py b/generative_ai/rag/retrieval_query_example.py index 6d949b8268b..7ce9e388558 100644 --- a/generative_ai/rag/retrieval_query_example.py +++ b/generative_ai/rag/retrieval_query_example.py @@ -24,29 +24,38 @@ def retrieval_query( ) -> RetrieveContextsResponse: # [START generativeaionvertexai_rag_retrieval_query] - from vertexai import rag - import vertexai + import agentplatform + + from agentplatform import types + from google import genai + from google.genai import types as genai_types # 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") + # Initialize Agent Platform client once per session + client = agentplatform.Client(project=PROJECT_ID, location="us-east4") - 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!", - rag_retrieval_config=rag.RagRetrievalConfig( - top_k=10, - filter=rag.utils.resources.Filter(vector_distance_threshold=0.5), + response = client.rag.retrieve_contexts( + vertex_rag_store=genai_types.VertexRagStore( + rag_resources=[ + genai_types.VertexRagStoreRagResource( + rag_corpus=corpus_name, + # Optional: supply IDs from `rag.list_files()`. + # rag_file_ids=["rag-file-1", "rag-file-2", ...], + ) + ], ), + query=types.RagQuery( + text="Hello World!", + rag_retrieval_config=genai_types.RagRetrievalConfig( + top_k=10, + filter=genai_types.RagRetrievalConfigFilter( + vector_distance_threshold=0.5 + ), + ), + ) ) print(response) # Example response: diff --git a/generative_ai/rag/upload_file_example.py b/generative_ai/rag/upload_file_example.py index f56cf23f2dc..51748e279b3 100644 --- a/generative_ai/rag/upload_file_example.py +++ b/generative_ai/rag/upload_file_example.py @@ -29,28 +29,25 @@ def upload_file( ) -> rag.RagFile: # [START generativeaionvertexai_rag_upload_file] - from vertexai import rag - import vertexai + import agentplatform # 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") + # Initialize Agent Platform client once per session + client = agentplatform.Client(project=PROJECT_ID, location="us-east4") - rag_file = rag.upload_file( + rag_file = client.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') + # display_name='file_display_name') # [END generativeaionvertexai_rag_upload_file] return rag_file From a8dfedd57f4e6b42fad5a5bd2cea554bc855e502 Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Wed, 1 Jul 2026 16:49:49 -0400 Subject: [PATCH 05/15] update additional rag samples --- generative_ai/rag/create_corpus_example.py | 4 +- .../create_corpus_feature_store_example.py | 1 + .../rag/create_corpus_pinecone_example.py | 38 +++++++++---------- .../create_corpus_vector_search_example.py | 37 +++++++++--------- .../create_corpus_vertex_ai_search_example.py | 22 +++++------ .../rag/create_corpus_weaviate_example.py | 35 +++++++++-------- 6 files changed, 69 insertions(+), 68 deletions(-) diff --git a/generative_ai/rag/create_corpus_example.py b/generative_ai/rag/create_corpus_example.py index c68801bc347..88e859db84d 100644 --- a/generative_ai/rag/create_corpus_example.py +++ b/generative_ai/rag/create_corpus_example.py @@ -44,15 +44,13 @@ def create_corpus( ) ) ) - client.rag.update_config( - updated_config=backend_config - ) # Create a corpus corpus = client.rag.create_corpus( rag_corpus=types.RagCorpus( display_name=display_name, description=description, + rag_vector_db_config=backend_config, ) ) print(corpus) diff --git a/generative_ai/rag/create_corpus_feature_store_example.py b/generative_ai/rag/create_corpus_feature_store_example.py index b0422f009d2..1020fc5eae0 100644 --- a/generative_ai/rag/create_corpus_feature_store_example.py +++ b/generative_ai/rag/create_corpus_feature_store_example.py @@ -58,6 +58,7 @@ def create_corpus_feature_store( rag_corpus=types.RagCorpus( display_name=display_name, description=description, + rag_vector_db_config=backend_config, ) ) print(corpus) diff --git a/generative_ai/rag/create_corpus_pinecone_example.py b/generative_ai/rag/create_corpus_pinecone_example.py index ebca30385e8..1dfcdd9b1ac 100644 --- a/generative_ai/rag/create_corpus_pinecone_example.py +++ b/generative_ai/rag/create_corpus_pinecone_example.py @@ -15,8 +15,6 @@ from typing import Optional -from vertexai.preview.rag import RagCorpus - PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") @@ -28,39 +26,39 @@ def create_corpus_pinecone( ) -> RagCorpus: # [START generativeaionvertexai_rag_create_corpus_pinecone] - from vertexai import rag - import vertexai + import agentplatform + from agentplatform import types # TODO(developer): Update and un-comment below lines # PROJECT_ID = "your-project-id" # pinecone_index_name = "pinecone-index-name" - # pinecone_api_key_secret_manager_version = "projects/{PROJECT_ID}/secrets/{SECRET_NAME}/versions/latest" # display_name = "test_corpus" # description = "Corpus Description" - # Initialize Vertex AI API once per session - vertexai.init(project=PROJECT_ID, location="us-central1") + # Initialize Agent Platform client once per session + client = agentplatform.Client(project=PROJECT_ID, location="us-central1") # Configure embedding model (Optional) - embedding_model_config = rag.RagEmbeddingModelConfig( - vertex_prediction_endpoint=rag.VertexPredictionEndpoint( - publisher_model="publishers/google/models/text-embedding-005" + embedding_model_config = types.RagEmbeddingModelConfig( + vertex_prediction_endpoint=types.RagEmbeddingModelConfigVertexPredictionEndpoint( + endpoint="publishers/google/models/text-embedding-005" ) ) # Configure Vector DB - vector_db = rag.Pinecone( - index_name=pinecone_index_name, - api_key=pinecone_api_key_secret_manager_version, + vector_db = types.RagVectorDbConfig( + pinecone=types.RagVectorDbConfigPinecone( + index_name=pinecone_index_name, + ), + rag_embedding_model_config=embedding_model_config, ) - corpus = rag.create_corpus( - display_name=display_name, - description=description, - backend_config=rag.RagVectorDbConfig( - rag_embedding_model_config=embedding_model_config, - vector_db=vector_db, - ), + corpus = client.rag.create_corpus( + rag_corpus=types.RagCorpus( + display_name=display_name, + description=description, + rag_vector_db_config=vector_db, + ) ) print(corpus) # Example response: diff --git a/generative_ai/rag/create_corpus_vector_search_example.py b/generative_ai/rag/create_corpus_vector_search_example.py index 5db30008046..f87554b823d 100644 --- a/generative_ai/rag/create_corpus_vector_search_example.py +++ b/generative_ai/rag/create_corpus_vector_search_example.py @@ -15,8 +15,6 @@ from typing import Optional -from vertexai.preview.rag import RagCorpus - PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") @@ -28,8 +26,8 @@ def create_corpus_vector_search( ) -> RagCorpus: # [START generativeaionvertexai_rag_create_corpus_vector_search] - from vertexai import rag - import vertexai + import agentplatform + from agentplatform import types # TODO(developer): Update and un-comment below lines # PROJECT_ID = "your-project-id" @@ -38,28 +36,31 @@ def create_corpus_vector_search( # display_name = "test_corpus" # description = "Corpus Description" - # Initialize Vertex AI API once per session - vertexai.init(project=PROJECT_ID, location="us-central1") + # Initialize Agent Platform client once per session + client = agentplatform.Client(project=PROJECT_ID, location="us-central1") # Configure embedding model (Optional) - embedding_model_config = rag.RagEmbeddingModelConfig( - vertex_prediction_endpoint=rag.VertexPredictionEndpoint( - publisher_model="publishers/google/models/text-embedding-005" + embedding_model_config = types.RagEmbeddingModelConfig( + vertex_prediction_endpoint=types.RagEmbeddingModelConfigVertexPredictionEndpoint( + endpoint="publishers/google/models/text-embedding-005" ) ) # Configure Vector DB - vector_db = rag.VertexVectorSearch( - index=vector_search_index_name, index_endpoint=vector_search_index_endpoint_name + vector_db = types.RagVectorDbConfigVertexVectorSearch( + index=vector_search_index_name, + index_endpoint=vector_search_index_endpoint_name ) - corpus = rag.create_corpus( - display_name=display_name, - description=description, - backend_config=rag.RagVectorDbConfig( - rag_embedding_model_config=embedding_model_config, - vector_db=vector_db, - ), + corpus = client.rag.create_corpus( + rag_corpus=types.RagCorpus( + display_name=display_name, + description=description, + rag_vector_db_config=types.RagVectorDbConfig( + rag_embedding_model_config=embedding_model_config, + vertex_vector_search=vector_db, + ), + ) ) print(corpus) # Example response: diff --git a/generative_ai/rag/create_corpus_vertex_ai_search_example.py b/generative_ai/rag/create_corpus_vertex_ai_search_example.py index 6d3fca5ab9c..8ac6cb3cb03 100644 --- a/generative_ai/rag/create_corpus_vertex_ai_search_example.py +++ b/generative_ai/rag/create_corpus_vertex_ai_search_example.py @@ -15,8 +15,6 @@ from typing import Optional -from vertexai import rag - PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") @@ -27,8 +25,8 @@ def create_corpus_vertex_ai_search( ) -> rag.RagCorpus: # [START generativeaionvertexai_rag_create_corpus_vertex_ai_search] - from vertexai import rag - import vertexai + import agentplatform + from agentplatform import types # TODO(developer): Update and un-comment below lines # PROJECT_ID = "your-project-id" @@ -36,18 +34,20 @@ def create_corpus_vertex_ai_search( # display_name = "test_corpus" # description = "Corpus Description" - # Initialize Vertex AI API once per session - vertexai.init(project=PROJECT_ID, location="us-central1") + # Initialize Agent Platform client once per session + client = agentplatform.Client(project=PROJECT_ID, location="us-central1") # Configure Search - vertex_ai_search_config = rag.VertexAiSearchConfig( + vertex_ai_search_config = types.VertexAiSearchConfig( serving_config=f"{vertex_ai_search_engine_name}/servingConfigs/default_search", ) - corpus = rag.create_corpus( - display_name=display_name, - description=description, - vertex_ai_search_config=vertex_ai_search_config, + corpus = client.rag.create_corpus( + rag_corpus=types.RagCorpus( + display_name=display_name, + description=description, + vertex_ai_search_config=vertex_ai_search_config, + ), ) print(corpus) # Example response: diff --git a/generative_ai/rag/create_corpus_weaviate_example.py b/generative_ai/rag/create_corpus_weaviate_example.py index 9823b8332f8..d55a32f8f5b 100644 --- a/generative_ai/rag/create_corpus_weaviate_example.py +++ b/generative_ai/rag/create_corpus_weaviate_example.py @@ -15,8 +15,6 @@ from typing import Optional -from vertexai.preview.rag import RagCorpus - PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") @@ -29,8 +27,8 @@ def create_corpus_weaviate( ) -> RagCorpus: # [START generativeaionvertexai_rag_create_corpus_weaviate] - from vertexai.preview import rag - import vertexai + import agentplatform + from agentplatform import types # TODO(developer): Update and un-comment below lines # PROJECT_ID = "your-project-id" @@ -40,26 +38,31 @@ def create_corpus_weaviate( # display_name = "test_corpus" # description = "Corpus Description" - # Initialize Vertex AI API once per session - vertexai.init(project=PROJECT_ID, location="us-central1") + # Initialize Agent Platform client once per session + client = agentplatform.Client(project=PROJECT_ID, location="us-central1") # Configure embedding model (Optional) - embedding_model_config = rag.EmbeddingModelConfig( - publisher_model="publishers/google/models/text-embedding-004" + embedding_model_config = types.RagEmbeddingModelConfig( + vertex_prediction_endpoint=types.RagEmbeddingModelConfigVertexPredictionEndpoint( + endpoint="publishers/google/models/text-embedding-004" + ) ) # Configure Vector DB - vector_db = rag.Weaviate( - weaviate_http_endpoint=weaviate_http_endpoint, + vector_db = types.RagVectorDbConfigWeaviate( + http_endpoint=weaviate_http_endpoint, collection_name=weaviate_collection_name, - api_key=weaviate_api_key_secret_manager_version, ) - corpus = rag.create_corpus( - display_name=display_name, - description=description, - embedding_model_config=embedding_model_config, - vector_db=vector_db, + corpus = client.rag.create_corpus( + rag_corpus=types.RagCorpus( + display_name=display_name, + description=description, + rag_embedding_model_config=embedding_model_config, + vector_db=types.RagVectorDbConfig( + weaviate=vector_db + ), + ) ) print(corpus) # Example response: From 7bdf192ad274d952a6b840e1b71332ddd0f2e12b Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Wed, 1 Jul 2026 17:01:11 -0400 Subject: [PATCH 06/15] Fix sample typos --- .../rag/create_corpus_feature_store_example.py | 5 ----- generative_ai/rag/create_corpus_weaviate_example.py | 2 +- generative_ai/rag/generate_content_example.py | 5 +---- generative_ai/rag/quickstart_example.py | 10 ++++++++++ 4 files changed, 12 insertions(+), 10 deletions(-) diff --git a/generative_ai/rag/create_corpus_feature_store_example.py b/generative_ai/rag/create_corpus_feature_store_example.py index 1020fc5eae0..77d86df5e76 100644 --- a/generative_ai/rag/create_corpus_feature_store_example.py +++ b/generative_ai/rag/create_corpus_feature_store_example.py @@ -49,11 +49,6 @@ def create_corpus_feature_store( ) ) - # Configure Vector DB - client.rag.update_config( - updated_config=backend_config - ) - corpus = client.rag.create_corpus( rag_corpus=types.RagCorpus( display_name=display_name, diff --git a/generative_ai/rag/create_corpus_weaviate_example.py b/generative_ai/rag/create_corpus_weaviate_example.py index d55a32f8f5b..e90929f9c85 100644 --- a/generative_ai/rag/create_corpus_weaviate_example.py +++ b/generative_ai/rag/create_corpus_weaviate_example.py @@ -59,7 +59,7 @@ def create_corpus_weaviate( display_name=display_name, description=description, rag_embedding_model_config=embedding_model_config, - vector_db=types.RagVectorDbConfig( + rag_vector_db_config=types.RagVectorDbConfig( weaviate=vector_db ), ) diff --git a/generative_ai/rag/generate_content_example.py b/generative_ai/rag/generate_content_example.py index 4fde2cb2256..ee5407f76f7 100644 --- a/generative_ai/rag/generate_content_example.py +++ b/generative_ai/rag/generate_content_example.py @@ -54,10 +54,7 @@ def generate_content_with_rag( # Create a GenAI SDK client to make a generate_content request genai_client = genai.Client(enterprise=True, project=PROJECT_ID, location="us-central1") - rag_model = GenerativeModel( - model_name="gemini-2.0-flash-001", tools=[rag_retrieval_tool] - ) - response = client.models.generate_content( + response = genai_client.models.generate_content( model="gemini-2.5-pro", contents="Why is the sky blue?", config=genai_types.GenerateContentConfig( diff --git a/generative_ai/rag/quickstart_example.py b/generative_ai/rag/quickstart_example.py index f8fcf6af85c..d7faef57353 100644 --- a/generative_ai/rag/quickstart_example.py +++ b/generative_ai/rag/quickstart_example.py @@ -48,10 +48,20 @@ def quickstart( # Initialize Agent Platform client once per session client = agentplatform.Client(project=PROJECT_ID, location="us-east4") + # Configure embedding model, for example "text-embedding-005". + embedding_model_config = types.RagEmbeddingModelConfig( + vertex_prediction_endpoint=types.RagEmbeddingModelConfigVertexPredictionEndpoint( + endpoint="publishers/google/models/text-embedding-005" + ), + ) + # Create RagCorpus rag_corpus = client.rag.create_corpus( rag_corpus=types.RagCorpus( display_name=display_name, + rag_vector_db_config=types.RagVectorDbConfig( + rag_embedding_model_config=embedding_model_config + ) ) ) From 146e150df8f55a657e48caa2749c4ca3d5c32327 Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Mon, 6 Jul 2026 10:49:42 -0400 Subject: [PATCH 07/15] Update imports and return types in RAG samples --- generative_ai/rag/create_corpus_example.py | 4 +++- generative_ai/rag/create_corpus_feature_store_example.py | 5 ++++- generative_ai/rag/create_corpus_pinecone_example.py | 4 +++- generative_ai/rag/create_corpus_vector_search_example.py | 4 +++- .../rag/create_corpus_vertex_ai_search_example.py | 5 ++++- generative_ai/rag/create_corpus_weaviate_example.py | 5 ++++- generative_ai/rag/generate_content_example.py | 5 ++++- generative_ai/rag/get_corpus_example.py | 5 +++-- generative_ai/rag/get_file_example.py | 5 +++-- generative_ai/rag/import_files_async_example.py | 5 +++-- generative_ai/rag/import_files_example.py | 5 +++-- generative_ai/rag/list_corpora_example.py | 7 +++---- generative_ai/rag/list_files_example.py | 7 +++---- generative_ai/rag/quickstart_example.py | 2 +- generative_ai/rag/retrieval_query_example.py | 5 +++-- generative_ai/rag/upload_file_example.py | 5 +++-- 16 files changed, 50 insertions(+), 28 deletions(-) diff --git a/generative_ai/rag/create_corpus_example.py b/generative_ai/rag/create_corpus_example.py index 88e859db84d..c9f4a43aab5 100644 --- a/generative_ai/rag/create_corpus_example.py +++ b/generative_ai/rag/create_corpus_example.py @@ -15,6 +15,8 @@ from typing import Optional +import agentplatform +from agentplatform import types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") @@ -22,7 +24,7 @@ def create_corpus( display_name: Optional[str] = None, description: Optional[str] = None, -) -> RagCorpus: +) -> types.RagCorpus: # [START generativeaionvertexai_rag_create_corpus] import agentplatform diff --git a/generative_ai/rag/create_corpus_feature_store_example.py b/generative_ai/rag/create_corpus_feature_store_example.py index 77d86df5e76..ffd96787dd0 100644 --- a/generative_ai/rag/create_corpus_feature_store_example.py +++ b/generative_ai/rag/create_corpus_feature_store_example.py @@ -15,6 +15,9 @@ from typing import Optional +import agentplatform +from agentplatform import types + PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") @@ -22,7 +25,7 @@ def create_corpus_feature_store( feature_view_name: str, display_name: Optional[str] = None, description: Optional[str] = None, -) -> RagCorpus: +) -> types.RagCorpus: # [START generativeaionvertexai_rag_create_corpus_feature_store] import agentplatform diff --git a/generative_ai/rag/create_corpus_pinecone_example.py b/generative_ai/rag/create_corpus_pinecone_example.py index 1dfcdd9b1ac..db26032f92b 100644 --- a/generative_ai/rag/create_corpus_pinecone_example.py +++ b/generative_ai/rag/create_corpus_pinecone_example.py @@ -17,13 +17,15 @@ PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") +import agentplatform +from agentplatform import types def create_corpus_pinecone( pinecone_index_name: str, pinecone_api_key_secret_manager_version: str, display_name: Optional[str] = None, description: Optional[str] = None, -) -> RagCorpus: +) -> types.RagCorpus: # [START generativeaionvertexai_rag_create_corpus_pinecone] import agentplatform diff --git a/generative_ai/rag/create_corpus_vector_search_example.py b/generative_ai/rag/create_corpus_vector_search_example.py index f87554b823d..0df0bb7fe72 100644 --- a/generative_ai/rag/create_corpus_vector_search_example.py +++ b/generative_ai/rag/create_corpus_vector_search_example.py @@ -17,13 +17,15 @@ PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") +import agentplatform +from agentplatform import types def create_corpus_vector_search( vector_search_index_name: str, vector_search_index_endpoint_name: str, display_name: Optional[str] = None, description: Optional[str] = None, -) -> RagCorpus: +) -> types.RagCorpus: # [START generativeaionvertexai_rag_create_corpus_vector_search] import agentplatform diff --git a/generative_ai/rag/create_corpus_vertex_ai_search_example.py b/generative_ai/rag/create_corpus_vertex_ai_search_example.py index 8ac6cb3cb03..6c0de292f92 100644 --- a/generative_ai/rag/create_corpus_vertex_ai_search_example.py +++ b/generative_ai/rag/create_corpus_vertex_ai_search_example.py @@ -15,6 +15,9 @@ from typing import Optional +import agentplatform +from agentplatform import types + PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") @@ -22,7 +25,7 @@ def create_corpus_vertex_ai_search( vertex_ai_search_engine_name: str, display_name: Optional[str] = None, description: Optional[str] = None, -) -> rag.RagCorpus: +) -> types.RagCorpus: # [START generativeaionvertexai_rag_create_corpus_vertex_ai_search] import agentplatform diff --git a/generative_ai/rag/create_corpus_weaviate_example.py b/generative_ai/rag/create_corpus_weaviate_example.py index e90929f9c85..fbb7127e41c 100644 --- a/generative_ai/rag/create_corpus_weaviate_example.py +++ b/generative_ai/rag/create_corpus_weaviate_example.py @@ -15,6 +15,9 @@ from typing import Optional +import agentplatform +from agentplatform import types + PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") @@ -24,7 +27,7 @@ def create_corpus_weaviate( weaviate_api_key_secret_manager_version: str, display_name: Optional[str] = None, description: Optional[str] = None, -) -> RagCorpus: +) -> types.RagCorpus: # [START generativeaionvertexai_rag_create_corpus_weaviate] import agentplatform diff --git a/generative_ai/rag/generate_content_example.py b/generative_ai/rag/generate_content_example.py index ee5407f76f7..196ae2e451b 100644 --- a/generative_ai/rag/generate_content_example.py +++ b/generative_ai/rag/generate_content_example.py @@ -14,12 +14,15 @@ import os +from google import genai +from genai import types as genai_types + PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") def generate_content_with_rag( corpus_name: str, -) -> GenerationResponse: +) -> genai_types.GenerateContentResponse: # [START generativeaionvertexai_rag_generate_content] import agentplatform diff --git a/generative_ai/rag/get_corpus_example.py b/generative_ai/rag/get_corpus_example.py index 17037287d5c..dca421f15d1 100644 --- a/generative_ai/rag/get_corpus_example.py +++ b/generative_ai/rag/get_corpus_example.py @@ -14,12 +14,13 @@ import os -from google.cloud.aiplatform_v1beta1 import RagCorpus +import agentplatform +from agentplatform import types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") -def get_corpus(corpus_name: str) -> RagCorpus: +def get_corpus(corpus_name: str) -> types.RagCorpus: # [START generativeaionvertexai_rag_get_corpus] import agentplatform diff --git a/generative_ai/rag/get_file_example.py b/generative_ai/rag/get_file_example.py index afdfa7ce625..73eb4b7c206 100644 --- a/generative_ai/rag/get_file_example.py +++ b/generative_ai/rag/get_file_example.py @@ -14,12 +14,13 @@ import os -from google.cloud.aiplatform_v1 import RagFile +import agentplatform +from agentplatform import types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") -def get_file(file_name: str) -> RagFile: +def get_file(file_name: str) -> types.RagFile: # [START generativeaionvertexai_rag_get_file] import agentplatform diff --git a/generative_ai/rag/import_files_async_example.py b/generative_ai/rag/import_files_async_example.py index 9b00e0482cb..754cfb4f59c 100644 --- a/generative_ai/rag/import_files_async_example.py +++ b/generative_ai/rag/import_files_async_example.py @@ -16,7 +16,8 @@ from typing import List -from google.cloud.aiplatform_v1 import ImportRagFilesResponse +import agentplatform +from agentplatform import types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") @@ -24,7 +25,7 @@ async def import_files_async( corpus_name: str, paths: List[str], -) -> ImportRagFilesResponse: +) -> types.ImportRagFilesResponse: # [START generativeaionvertexai_rag_import_files_async] import agentplatform diff --git a/generative_ai/rag/import_files_example.py b/generative_ai/rag/import_files_example.py index 36bbc08af9b..b5984529f69 100644 --- a/generative_ai/rag/import_files_example.py +++ b/generative_ai/rag/import_files_example.py @@ -15,7 +15,8 @@ import os from typing import List -from google.cloud.aiplatform_v1 import ImportRagFilesResponse +import agentplatform +from agentplatform import types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") @@ -23,7 +24,7 @@ def import_files( corpus_name: str, paths: List[str], -) -> ImportRagFilesResponse: +) -> types.ImportRagFilesResponse: # [START generativeaionvertexai_rag_import_files] import agentplatform diff --git a/generative_ai/rag/list_corpora_example.py b/generative_ai/rag/list_corpora_example.py index 9bdab807b5f..1af229ee3c3 100644 --- a/generative_ai/rag/list_corpora_example.py +++ b/generative_ai/rag/list_corpora_example.py @@ -14,14 +14,13 @@ import os -from google.cloud.aiplatform_v1beta1.services.vertex_rag_data_service.pagers import ( - ListRagCorporaPager, -) +import agentplatform +from agentplatform import types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") -def list_corpora() -> ListRagCorporaPager: +def list_corpora() -> types.ListRagCorporaResponse: # [START generativeaionvertexai_rag_list_corpora] import agentplatform diff --git a/generative_ai/rag/list_files_example.py b/generative_ai/rag/list_files_example.py index ce9e8aeb19e..0dc408c4fbd 100644 --- a/generative_ai/rag/list_files_example.py +++ b/generative_ai/rag/list_files_example.py @@ -14,14 +14,13 @@ import os -from google.cloud.aiplatform_v1.services.vertex_rag_data_service.pagers import ( - ListRagFilesPager, -) +import agentplatform +from agentplatform import types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") -def list_files(corpus_name: str) -> ListRagFilesPager: +def list_files(corpus_name: str) -> types.ListRagFilesResponse: # [START generativeaionvertexai_rag_list_files] import agentplatform diff --git a/generative_ai/rag/quickstart_example.py b/generative_ai/rag/quickstart_example.py index d7faef57353..343257f5acf 100644 --- a/generative_ai/rag/quickstart_example.py +++ b/generative_ai/rag/quickstart_example.py @@ -27,7 +27,7 @@ def quickstart( display_name: str, paths: List[str], -) -> Tuple[rag.RagCorpus, GenerationResponse]: +) -> Tuple[types.RagCorpus, genai_types.GenerateContentResponse]: # [START generativeaionvertexai_rag_quickstart] import agentplatform diff --git a/generative_ai/rag/retrieval_query_example.py b/generative_ai/rag/retrieval_query_example.py index 7ce9e388558..61aa28551ad 100644 --- a/generative_ai/rag/retrieval_query_example.py +++ b/generative_ai/rag/retrieval_query_example.py @@ -14,14 +14,15 @@ import os -from google.cloud.aiplatform_v1beta1 import RetrieveContextsResponse +import agentplatform +from agentplatform import types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") def retrieval_query( corpus_name: str, -) -> RetrieveContextsResponse: +) -> types.RetrieveContextsResponse: # [START generativeaionvertexai_rag_retrieval_query] import agentplatform diff --git a/generative_ai/rag/upload_file_example.py b/generative_ai/rag/upload_file_example.py index 51748e279b3..1cadef7f717 100644 --- a/generative_ai/rag/upload_file_example.py +++ b/generative_ai/rag/upload_file_example.py @@ -16,7 +16,8 @@ from typing import Optional -from vertexai import rag +import agentplatform +from agentplatform import types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") @@ -26,7 +27,7 @@ def upload_file( path: str, display_name: Optional[str] = None, description: Optional[str] = None, -) -> rag.RagFile: +) -> types.RagFile: # [START generativeaionvertexai_rag_upload_file] import agentplatform From 43524dda1c7a661e5b2f33ceab5b1d167c195bc8 Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Mon, 6 Jul 2026 11:07:49 -0400 Subject: [PATCH 08/15] Remove unused imports --- generative_ai/rag/create_corpus_example.py | 1 - generative_ai/rag/create_corpus_feature_store_example.py | 1 - generative_ai/rag/create_corpus_pinecone_example.py | 1 - generative_ai/rag/create_corpus_vector_search_example.py | 1 - generative_ai/rag/create_corpus_vertex_ai_search_example.py | 1 - generative_ai/rag/create_corpus_weaviate_example.py | 1 - generative_ai/rag/generate_content_example.py | 1 - generative_ai/rag/get_corpus_example.py | 1 - generative_ai/rag/get_file_example.py | 1 - generative_ai/rag/import_files_async_example.py | 1 - generative_ai/rag/import_files_example.py | 1 - generative_ai/rag/list_corpora_example.py | 1 - generative_ai/rag/list_files_example.py | 1 - generative_ai/rag/quickstart_example.py | 2 -- generative_ai/rag/retrieval_query_example.py | 1 - generative_ai/rag/upload_file_example.py | 1 - 16 files changed, 17 deletions(-) diff --git a/generative_ai/rag/create_corpus_example.py b/generative_ai/rag/create_corpus_example.py index c9f4a43aab5..f5bec091527 100644 --- a/generative_ai/rag/create_corpus_example.py +++ b/generative_ai/rag/create_corpus_example.py @@ -15,7 +15,6 @@ from typing import Optional -import agentplatform from agentplatform import types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") diff --git a/generative_ai/rag/create_corpus_feature_store_example.py b/generative_ai/rag/create_corpus_feature_store_example.py index ffd96787dd0..0295ea44420 100644 --- a/generative_ai/rag/create_corpus_feature_store_example.py +++ b/generative_ai/rag/create_corpus_feature_store_example.py @@ -15,7 +15,6 @@ from typing import Optional -import agentplatform from agentplatform import types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") diff --git a/generative_ai/rag/create_corpus_pinecone_example.py b/generative_ai/rag/create_corpus_pinecone_example.py index db26032f92b..031578b7d55 100644 --- a/generative_ai/rag/create_corpus_pinecone_example.py +++ b/generative_ai/rag/create_corpus_pinecone_example.py @@ -17,7 +17,6 @@ PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") -import agentplatform from agentplatform import types def create_corpus_pinecone( diff --git a/generative_ai/rag/create_corpus_vector_search_example.py b/generative_ai/rag/create_corpus_vector_search_example.py index 0df0bb7fe72..b5d7b0ab671 100644 --- a/generative_ai/rag/create_corpus_vector_search_example.py +++ b/generative_ai/rag/create_corpus_vector_search_example.py @@ -17,7 +17,6 @@ PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") -import agentplatform from agentplatform import types def create_corpus_vector_search( diff --git a/generative_ai/rag/create_corpus_vertex_ai_search_example.py b/generative_ai/rag/create_corpus_vertex_ai_search_example.py index 6c0de292f92..01231e73a16 100644 --- a/generative_ai/rag/create_corpus_vertex_ai_search_example.py +++ b/generative_ai/rag/create_corpus_vertex_ai_search_example.py @@ -15,7 +15,6 @@ from typing import Optional -import agentplatform from agentplatform import types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") diff --git a/generative_ai/rag/create_corpus_weaviate_example.py b/generative_ai/rag/create_corpus_weaviate_example.py index fbb7127e41c..a12452af2e6 100644 --- a/generative_ai/rag/create_corpus_weaviate_example.py +++ b/generative_ai/rag/create_corpus_weaviate_example.py @@ -15,7 +15,6 @@ from typing import Optional -import agentplatform from agentplatform import types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") diff --git a/generative_ai/rag/generate_content_example.py b/generative_ai/rag/generate_content_example.py index 196ae2e451b..2867d3eca31 100644 --- a/generative_ai/rag/generate_content_example.py +++ b/generative_ai/rag/generate_content_example.py @@ -14,7 +14,6 @@ import os -from google import genai from genai import types as genai_types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") diff --git a/generative_ai/rag/get_corpus_example.py b/generative_ai/rag/get_corpus_example.py index dca421f15d1..2e0a1ae868c 100644 --- a/generative_ai/rag/get_corpus_example.py +++ b/generative_ai/rag/get_corpus_example.py @@ -14,7 +14,6 @@ import os -import agentplatform from agentplatform import types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") diff --git a/generative_ai/rag/get_file_example.py b/generative_ai/rag/get_file_example.py index 73eb4b7c206..af64b386f83 100644 --- a/generative_ai/rag/get_file_example.py +++ b/generative_ai/rag/get_file_example.py @@ -14,7 +14,6 @@ import os -import agentplatform from agentplatform import types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") diff --git a/generative_ai/rag/import_files_async_example.py b/generative_ai/rag/import_files_async_example.py index 754cfb4f59c..35a274f3e0e 100644 --- a/generative_ai/rag/import_files_async_example.py +++ b/generative_ai/rag/import_files_async_example.py @@ -16,7 +16,6 @@ from typing import List -import agentplatform from agentplatform import types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") diff --git a/generative_ai/rag/import_files_example.py b/generative_ai/rag/import_files_example.py index b5984529f69..d2d92192c8a 100644 --- a/generative_ai/rag/import_files_example.py +++ b/generative_ai/rag/import_files_example.py @@ -15,7 +15,6 @@ import os from typing import List -import agentplatform from agentplatform import types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") diff --git a/generative_ai/rag/list_corpora_example.py b/generative_ai/rag/list_corpora_example.py index 1af229ee3c3..1e9ba774384 100644 --- a/generative_ai/rag/list_corpora_example.py +++ b/generative_ai/rag/list_corpora_example.py @@ -14,7 +14,6 @@ import os -import agentplatform from agentplatform import types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") diff --git a/generative_ai/rag/list_files_example.py b/generative_ai/rag/list_files_example.py index 0dc408c4fbd..f875833c081 100644 --- a/generative_ai/rag/list_files_example.py +++ b/generative_ai/rag/list_files_example.py @@ -14,7 +14,6 @@ import os -import agentplatform from agentplatform import types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") diff --git a/generative_ai/rag/quickstart_example.py b/generative_ai/rag/quickstart_example.py index 343257f5acf..5489110a67e 100644 --- a/generative_ai/rag/quickstart_example.py +++ b/generative_ai/rag/quickstart_example.py @@ -13,12 +13,10 @@ # limitations under the License. import os -import agentplatform from typing import List, Tuple from agentplatform import types -from google import genai from google.genai import types as genai_types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") diff --git a/generative_ai/rag/retrieval_query_example.py b/generative_ai/rag/retrieval_query_example.py index 61aa28551ad..71f8b57abf0 100644 --- a/generative_ai/rag/retrieval_query_example.py +++ b/generative_ai/rag/retrieval_query_example.py @@ -14,7 +14,6 @@ import os -import agentplatform from agentplatform import types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") diff --git a/generative_ai/rag/upload_file_example.py b/generative_ai/rag/upload_file_example.py index 1cadef7f717..d3d964d700a 100644 --- a/generative_ai/rag/upload_file_example.py +++ b/generative_ai/rag/upload_file_example.py @@ -16,7 +16,6 @@ from typing import Optional -import agentplatform from agentplatform import types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") From 5081f6c225c32fd3168559037b53ea97e79a5f18 Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Mon, 6 Jul 2026 11:13:30 -0400 Subject: [PATCH 09/15] Fix linter errors --- generative_ai/rag/create_corpus_pinecone_example.py | 3 ++- generative_ai/rag/create_corpus_vector_search_example.py | 3 +-- generative_ai/rag/generate_content_example.py | 6 ------ generative_ai/rag/import_files_async_example.py | 2 +- generative_ai/rag/list_files_example.py | 2 +- generative_ai/rag/retrieval_query_example.py | 1 - 6 files changed, 5 insertions(+), 12 deletions(-) diff --git a/generative_ai/rag/create_corpus_pinecone_example.py b/generative_ai/rag/create_corpus_pinecone_example.py index 031578b7d55..d4cb83dfffb 100644 --- a/generative_ai/rag/create_corpus_pinecone_example.py +++ b/generative_ai/rag/create_corpus_pinecone_example.py @@ -13,11 +13,12 @@ # limitations under the License. import os +from agentplatform import types from typing import Optional PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") -from agentplatform import types + def create_corpus_pinecone( pinecone_index_name: str, diff --git a/generative_ai/rag/create_corpus_vector_search_example.py b/generative_ai/rag/create_corpus_vector_search_example.py index b5d7b0ab671..8bcd0836aaf 100644 --- a/generative_ai/rag/create_corpus_vector_search_example.py +++ b/generative_ai/rag/create_corpus_vector_search_example.py @@ -13,12 +13,11 @@ # limitations under the License. import os +from agentplatform import types from typing import Optional PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") -from agentplatform import types - def create_corpus_vector_search( vector_search_index_name: str, vector_search_index_endpoint_name: str, diff --git a/generative_ai/rag/generate_content_example.py b/generative_ai/rag/generate_content_example.py index 2867d3eca31..d770bf2eb2c 100644 --- a/generative_ai/rag/generate_content_example.py +++ b/generative_ai/rag/generate_content_example.py @@ -24,9 +24,6 @@ def generate_content_with_rag( ) -> genai_types.GenerateContentResponse: # [START generativeaionvertexai_rag_generate_content] - import agentplatform - - from agentplatform import types from google import genai from google.genai import types as genai_types @@ -34,9 +31,6 @@ def generate_content_with_rag( # PROJECT_ID = "your-project-id" # corpus_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}" - # Initialize Agent Platform client once per session - client = agentplatform.Client(project=PROJECT_ID, location="us-central1") - rag_retrieval_tool = genai_types.Tool( retrieval=genai_types.Retrieval( vertex_rag_store=genai_types.VertexRagStore( diff --git a/generative_ai/rag/import_files_async_example.py b/generative_ai/rag/import_files_async_example.py index 35a274f3e0e..8ca344ce9ee 100644 --- a/generative_ai/rag/import_files_async_example.py +++ b/generative_ai/rag/import_files_async_example.py @@ -69,7 +69,7 @@ async def import_files_async( # Imported 2 files. # [END generativeaionvertexai_rag_import_files_async] - return result + return response if __name__ == "__main__": diff --git a/generative_ai/rag/list_files_example.py b/generative_ai/rag/list_files_example.py index f875833c081..7da8e6f8ac5 100644 --- a/generative_ai/rag/list_files_example.py +++ b/generative_ai/rag/list_files_example.py @@ -42,7 +42,7 @@ def list_files(corpus_name: str) -> types.ListRagFilesResponse: # projects/1234567890/locations/us-central1/ragCorpora/111111111111/ragFiles/333333333333 # [END generativeaionvertexai_rag_list_files] - return files + return files_response if __name__ == "__main__": diff --git a/generative_ai/rag/retrieval_query_example.py b/generative_ai/rag/retrieval_query_example.py index 71f8b57abf0..c887298550b 100644 --- a/generative_ai/rag/retrieval_query_example.py +++ b/generative_ai/rag/retrieval_query_example.py @@ -27,7 +27,6 @@ def retrieval_query( import agentplatform from agentplatform import types - from google import genai from google.genai import types as genai_types # TODO(developer): Update and un-comment below lines From 42051c80bb0c44179ec3f7c01b3980b0ee6c4654 Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Mon, 6 Jul 2026 11:17:20 -0400 Subject: [PATCH 10/15] Fix linter errors --- generative_ai/rag/create_corpus_pinecone_example.py | 3 ++- generative_ai/rag/create_corpus_vector_search_example.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/generative_ai/rag/create_corpus_pinecone_example.py b/generative_ai/rag/create_corpus_pinecone_example.py index d4cb83dfffb..74a2c6081e5 100644 --- a/generative_ai/rag/create_corpus_pinecone_example.py +++ b/generative_ai/rag/create_corpus_pinecone_example.py @@ -13,9 +13,10 @@ # limitations under the License. import os -from agentplatform import types from typing import Optional +from agentplatform import types + PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") diff --git a/generative_ai/rag/create_corpus_vector_search_example.py b/generative_ai/rag/create_corpus_vector_search_example.py index 8bcd0836aaf..8f3905adbd1 100644 --- a/generative_ai/rag/create_corpus_vector_search_example.py +++ b/generative_ai/rag/create_corpus_vector_search_example.py @@ -13,9 +13,10 @@ # limitations under the License. import os -from agentplatform import types from typing import Optional +from agentplatform import types + PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") def create_corpus_vector_search( From ad3b42a6ef277d070d6cabc2dce2abecaec81e29 Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Mon, 6 Jul 2026 11:33:00 -0400 Subject: [PATCH 11/15] Fix more linter errors --- generative_ai/rag/quickstart_example.py | 11 ++--------- generative_ai/rag/test_rag_examples.py | 2 +- 2 files changed, 3 insertions(+), 10 deletions(-) diff --git a/generative_ai/rag/quickstart_example.py b/generative_ai/rag/quickstart_example.py index 5489110a67e..2d4440aeb33 100644 --- a/generative_ai/rag/quickstart_example.py +++ b/generative_ai/rag/quickstart_example.py @@ -20,11 +20,12 @@ from google.genai import types as genai_types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") +MODEL_ID = os.getenv("MODEL_ID") def quickstart( display_name: str, - paths: List[str], + gcs_path: str, ) -> Tuple[types.RagCorpus, genai_types.GenerateContentResponse]: # [START generativeaionvertexai_rag_quickstart] import agentplatform @@ -68,14 +69,6 @@ def quickstart( name=rag_corpus.name, import_config=types.ImportRagFilesConfig( gcs_source=genai_types.GcsSource(uris=[gcs_path]), - google_drive_source=types.GoogleDriveSource( - resource_ids=[ - types.GoogleDriveSourceResourceId( - resource_id=google_drive_path, - resource_type=types.ResourceType.RESOURCE_TYPE_FILE - ) - ] - ), # optional rag_file_transformation_config=types.RagFileTransformationConfig( rag_file_chunking_config=types.RagFileChunkingConfig( chunk_size=512, diff --git a/generative_ai/rag/test_rag_examples.py b/generative_ai/rag/test_rag_examples.py index 3d562f5463c..62d1275c791 100644 --- a/generative_ai/rag/test_rag_examples.py +++ b/generative_ai/rag/test_rag_examples.py @@ -210,7 +210,7 @@ def test_generate_content_with_rag(test_corpus: pytest.fixture) -> None: def test_quickstart() -> None: corpus, response = quickstart_example.quickstart( - "test_corpus_quickstart", [GCS_FILE] + "test_corpus_quickstart", GCS_FILE ) assert response assert response.text From db4cc3948c7910afca2a04c3ba6ce3cadf758bb0 Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Mon, 6 Jul 2026 11:56:45 -0400 Subject: [PATCH 12/15] Fix sample test runner --- generative_ai/rag/quickstart_example.py | 2 +- generative_ai/rag/test_rag_examples.py | 16 ++++++---------- 2 files changed, 7 insertions(+), 11 deletions(-) diff --git a/generative_ai/rag/quickstart_example.py b/generative_ai/rag/quickstart_example.py index 2d4440aeb33..b57ae049e9e 100644 --- a/generative_ai/rag/quickstart_example.py +++ b/generative_ai/rag/quickstart_example.py @@ -14,7 +14,7 @@ import os -from typing import List, Tuple +from typing import Tuple from agentplatform import types from google.genai import types as genai_types diff --git a/generative_ai/rag/test_rag_examples.py b/generative_ai/rag/test_rag_examples.py index 62d1275c791..715f1d6c92e 100644 --- a/generative_ai/rag/test_rag_examples.py +++ b/generative_ai/rag/test_rag_examples.py @@ -17,7 +17,6 @@ from pathlib import Path import pytest -import vertexai import create_corpus_example import create_corpus_feature_store_example @@ -47,9 +46,6 @@ GCS_FILE = "gs://cloud-samples-data/generative-ai/pdf/earnings_statement.pdf" -vertexai.init(project=PROJECT_ID, location=LOCATION) - - @pytest.fixture(scope="module", name="test_file") def test_file_fixture() -> None: file_path = Path("./hello.txt") @@ -156,14 +152,14 @@ def test_get_corpus(test_corpus: pytest.fixture) -> None: def test_list_corpora(test_corpus: pytest.fixture) -> None: corpora = list_corpora_example.list_corpora() - assert any(c.display_name == test_corpus.display_name for c in corpora) + assert any(c.display_name == test_corpus.display_name for c in corpora.rag_corpora) def test_upload_file(test_corpus: pytest.fixture, test_file: pytest.fixture) -> None: rag_file = upload_file_example.upload_file(test_corpus.name, test_file) assert rag_file files = list_files_example.list_files(test_corpus.name) - imported_file = next(iter(files)) + imported_file = next(iter(files.rag_files)) delete_file_example.delete_file(imported_file.name) @@ -171,7 +167,7 @@ def test_import_files(test_corpus: pytest.fixture) -> None: response = import_files_example.import_files(test_corpus.name, [GCS_FILE]) assert response.imported_rag_files_count > 0 files = list_files_example.list_files(test_corpus.name) - imported_file = next(iter(files)) + imported_file = next(iter(files.rag_files)) delete_file_example.delete_file(imported_file.name) @@ -182,7 +178,7 @@ async def test_import_files_async(test_corpus: pytest.fixture) -> None: ) assert result.imported_rag_files_count > 0 files = list_files_example.list_files(test_corpus.name) - imported_file = next(iter(files)) + imported_file = next(iter(files.rag_files)) delete_file_example.delete_file(imported_file.name) @@ -193,7 +189,7 @@ def test_get_file(uploaded_file: pytest.fixture) -> None: def test_list_files(test_corpus: pytest.fixture, uploaded_file: pytest.fixture) -> None: files = list_files_example.list_files(test_corpus.name) - assert any(f.name == uploaded_file.name for f in files) + assert any(f.name == uploaded_file.name for f in files.rag_files) def test_retrieval_query(test_corpus: pytest.fixture) -> None: @@ -214,4 +210,4 @@ def test_quickstart() -> None: ) assert response assert response.text - delete_corpus_example.delete_corpus(corpus.name) + delete_corpus_example.delete_corpus(corpus.name) \ No newline at end of file From b4eb0ca4e8fc056b91aba2a8c973b29976fd44d7 Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Mon, 6 Jul 2026 12:03:42 -0400 Subject: [PATCH 13/15] fix import order --- generative_ai/rag/test_rag_examples.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/generative_ai/rag/test_rag_examples.py b/generative_ai/rag/test_rag_examples.py index 715f1d6c92e..eb95f9e78d8 100644 --- a/generative_ai/rag/test_rag_examples.py +++ b/generative_ai/rag/test_rag_examples.py @@ -16,8 +16,6 @@ import os from pathlib import Path -import pytest - import create_corpus_example import create_corpus_feature_store_example import create_corpus_pinecone_example @@ -37,6 +35,8 @@ import retrieval_query_example import upload_file_example +import pytest + # TODO(https://github.com/GoogleCloudPlatform/python-docs-samples/issues/11557): Remove once Allowlist is removed pytest.skip(allow_module_level=True) From 39423e00f9da33b4f775f574946002179adae570 Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Mon, 6 Jul 2026 12:08:28 -0400 Subject: [PATCH 14/15] fix import error --- generative_ai/rag/test_rag_examples.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/generative_ai/rag/test_rag_examples.py b/generative_ai/rag/test_rag_examples.py index eb95f9e78d8..c35b3c63bba 100644 --- a/generative_ai/rag/test_rag_examples.py +++ b/generative_ai/rag/test_rag_examples.py @@ -31,12 +31,11 @@ import import_files_example import list_corpora_example import list_files_example +import pytest import quickstart_example import retrieval_query_example import upload_file_example -import pytest - # TODO(https://github.com/GoogleCloudPlatform/python-docs-samples/issues/11557): Remove once Allowlist is removed pytest.skip(allow_module_level=True) From 25ce4d0aa54a399a2b9c52ef91b89d2f028e5d83 Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Mon, 6 Jul 2026 12:59:38 -0400 Subject: [PATCH 15/15] Fix typos --- generative_ai/rag/create_corpus_feature_store_example.py | 2 +- generative_ai/rag/generate_content_example.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/generative_ai/rag/create_corpus_feature_store_example.py b/generative_ai/rag/create_corpus_feature_store_example.py index 0295ea44420..39f61898361 100644 --- a/generative_ai/rag/create_corpus_feature_store_example.py +++ b/generative_ai/rag/create_corpus_feature_store_example.py @@ -46,7 +46,7 @@ def create_corpus_feature_store( endpoint="publishers/google/models/text-embedding-005" ), ), - vertex_feature_store=types.RagDbConfigVertexFeatureStore( + vertex_feature_store=types.RagVectorDbConfigVertexFeatureStore( feature_view_resource_name=feature_view_name ) ) diff --git a/generative_ai/rag/generate_content_example.py b/generative_ai/rag/generate_content_example.py index d770bf2eb2c..6bff7cb1910 100644 --- a/generative_ai/rag/generate_content_example.py +++ b/generative_ai/rag/generate_content_example.py @@ -35,7 +35,9 @@ def generate_content_with_rag( retrieval=genai_types.Retrieval( vertex_rag_store=genai_types.VertexRagStore( rag_resources=[ - genai_types.VertexRagStoreRagResource + genai_types.VertexRagStoreRagResource( + rag_corpus=corpus_name + ) ], rag_retrieval_config=genai_types.RagRetrievalConfig( top_k=10,