Skip to content
This repository was archived by the owner on Aug 1, 2026. It is now read-only.

Commit f06becc

Browse files
sararobsararob
authored
feat: Update RAG notebook to use latest module (GoogleCloudPlatform#2962)
1 parent 0dd31ff commit f06becc

3 files changed

Lines changed: 81 additions & 120 deletions

File tree

‎.github/actions/spelling/allow.txt‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -335,6 +335,7 @@ Depatmint
335335
Derpanis
336336
descgen
337337
deskmates
338+
devkey
338339
Dexin
339340
dflags
340341
dfs
@@ -773,6 +774,7 @@ japonicum
773774
javac
774775
JAVACMD
775776
Jax
777+
jdbc
776778
JDK
777779
jdoe
778780
Jedi
@@ -886,6 +888,8 @@ linkedin
886888
linted
887889
linting
888890
litellm
891+
livekit
892+
LIVEKIT
889893
Livewire
890894
LLa
891895
Llion
@@ -1609,6 +1613,7 @@ tsne
16091613
TSNE
16101614
tsv
16111615
ttext
1616+
ttfb
16121617
TTFB
16131618
ttft
16141619
TTFT

‎.github/actions/spelling/expect.txt‎

Lines changed: 0 additions & 5 deletions
This file was deleted.

‎gemini/rag-engine/intro_rag_engine.ipynb‎

Lines changed: 76 additions & 115 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@
2929
"id": "JAPoU8Sm5E6e"
3030
},
3131
"source": [
32-
"# Intro to Building a Scalable and Modular RAG System with RAG Engine in Vertex AI \n",
32+
"# Intro to Building a Scalable and Modular RAG System with RAG Engine on Gemini Enterprise Agent Platform\n",
3333
"\n",
3434
"<table align=\"left\">\n",
3535
" <td style=\"text-align: center\">\n",
@@ -38,13 +38,13 @@
3838
" </a>\n",
3939
" </td>\n",
4040
" <td style=\"text-align: center\">\n",
41-
" <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Frag-engine%2Fintro_rag_engine.ipynb\">\n",
41+
" <a href=\"https://console.cloud.google.com/agent-platform/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Frag-engine%2Fintro_rag_engine.ipynb\">\n",
4242
" <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
4343
" </a>\n",
4444
" </td>\n",
4545
" <td style=\"text-align: center\">\n",
46-
" <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/rag-engine/intro_rag_engine.ipynb\">\n",
47-
" <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
46+
" <a href=\"https://console.cloud.google.com/agent-platform/workbench/instances?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/rag-engine/intro_rag_engine.ipynb\">\n",
47+
" <img width=\"32px\" src=\"https://storage.googleapis.com/github-repo/workbench-icon.svg\" alt=\"Workbench logo\"><br> Open in Workbench\n",
4848
" </a>\n",
4949
" </td>\n",
5050
" <td style=\"text-align: center\">\n",
@@ -56,6 +56,7 @@
5656
"\n",
5757
"<div style=\"clear: both;\"></div>\n",
5858
"\n",
59+
"<p>\n",
5960
"<b>Share to:</b>\n",
6061
"\n",
6162
"<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/intro_rag_engine.ipynb\" target=\"_blank\">\n",
@@ -76,7 +77,8 @@
7677
"\n",
7778
"<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/rag-engine/intro_rag_engine.ipynb\" target=\"_blank\">\n",
7879
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
79-
"</a> "
80+
"</a>\n",
81+
"</p> "
8082
]
8183
},
8284
{
@@ -109,7 +111,7 @@
109111
"have, a better context is provided. The improved context along with the query\n",
110112
"enhances the quality of the LLM's response.\n",
111113
"\n",
112-
"The following concepts are key to understanding Vertex AI RAG Engine. These concepts are listed in the order of the\n",
114+
"The following concepts are key to understanding Agent Platform RAG Engine. These concepts are listed in the order of the\n",
113115
"retrieval-augmented generation (RAG) process.\n",
114116
"\n",
115117
"1. **Data ingestion**: Intake data from different data sources. For example,\n",
@@ -135,7 +137,7 @@
135137
" original user query as a guide for the generative AI model to generate\n",
136138
" factually grounded and relevant responses.\n",
137139
"\n",
138-
"For more information, refer to the public documentation for [Vertex AI RAG Engine](https://cloud.google.com/vertex-ai/generative-ai/docs/rag-overview)."
140+
"For more information, refer to the public documentation for [Agent Platform RAG Engine](https://docs.cloud.google.com/gemini-enterprise-agent-platform/build/rag-engine/rag-overview)."
139141
]
140142
},
141143
{
@@ -153,7 +155,7 @@
153155
"id": "No17Cw5hgx12"
154156
},
155157
"source": [
156-
"### Install Vertex AI SDK and Google Gen AI SDK\n"
158+
"### Install Agent Platform SDK and Google Gen AI SDK\n"
157159
]
158160
},
159161
{
@@ -200,13 +202,13 @@
200202
"id": "DF4l8DTdWgPY"
201203
},
202204
"source": [
203-
"### Set Google Cloud project information and initialize Vertex AI SDK\n",
205+
"### Set Google Cloud project information\n",
204206
"\n",
205-
"To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
207+
"To get started using Agent Platform, you must have an existing Google Cloud project and [enable the Agent Platform API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
206208
"\n",
207-
"Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment).\n",
209+
"Learn more about [setting up a project](https://docs.cloud.google.com/resource-manager/docs/creating-managing-projects) and a [development environment](https://cloud.google.com/docs/authentication/set-up-adc-local-dev-environment).\n",
208210
"\n",
209-
"See [supported regions](https://cloud.google.com/vertex-ai/generative-ai/docs/rag-engine/rag-overview#supported-regions) for location options.\n"
211+
"See [supported regions](https://cloud.google.com/vertex-ai/generative-ai/docs/rag-engine/rag-overview#supported-regions) for location options."
210212
]
211213
},
212214
{
@@ -220,8 +222,7 @@
220222
"# Use the environment variable if the user doesn't provide Project ID.\n",
221223
"import os\n",
222224
"\n",
223-
"import vertexai\n",
224-
"from google import genai\n",
225+
"import agentplatform\n",
225226
"\n",
226227
"# fmt: off\n",
227228
"PROJECT_ID = \"[your-project-id]\" # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
@@ -230,8 +231,7 @@
230231
" PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
231232
"\n",
232233
"# See https://cloud.google.com/vertex-ai/generative-ai/docs/rag-engine/rag-overview#supported-regions for location options.\n",
233-
"vertexai.init(project=PROJECT_ID, location=\"us-east1\")\n",
234-
"client = genai.Client(vertexai=True, project=PROJECT_ID, location=\"global\")"
234+
"client = agentplatform.Client(project=PROJECT_ID, location=\"us-east1\")"
235235
]
236236
},
237237
{
@@ -252,8 +252,9 @@
252252
"outputs": [],
253253
"source": [
254254
"from IPython.display import Markdown, display\n",
255-
"from google.genai.types import GenerateContentConfig, Retrieval, Tool, VertexRagStore\n",
256-
"from vertexai import rag"
255+
"from agentplatform import types\n",
256+
"from google import genai\n",
257+
"from google.genai import types as genai_types"
257258
]
258259
},
259260
{
@@ -273,20 +274,10 @@
273274
},
274275
"outputs": [],
275276
"source": [
276-
"# Currently supports Google first-party embedding models\n",
277-
"# fmt: off\n",
278-
"EMBEDDING_MODEL = \"publishers/google/models/text-embedding-005\" # @param {type:\"string\", isTemplate: true}\n",
279-
"# fmt: on\n",
280-
"\n",
281-
"rag_corpus = rag.create_corpus(\n",
282-
" display_name=\"my-rag-corpus\",\n",
283-
" backend_config=rag.RagVectorDbConfig(\n",
284-
" rag_embedding_model_config=rag.RagEmbeddingModelConfig(\n",
285-
" vertex_prediction_endpoint=rag.VertexPredictionEndpoint(\n",
286-
" publisher_model=EMBEDDING_MODEL\n",
287-
" )\n",
288-
" )\n",
289-
" ),\n",
277+
"rag_corpus = client.rag.create_corpus(\n",
278+
" rag_corpus=types.RagCorpus(\n",
279+
" display_name=\"my-rag-corpus\",\n",
280+
" )\n",
290281
")"
291282
]
292283
},
@@ -307,7 +298,7 @@
307298
},
308299
"outputs": [],
309300
"source": [
310-
"rag.list_corpora()"
301+
"client.rag.list_corpora()"
311302
]
312303
},
313304
{
@@ -365,11 +356,10 @@
365356
},
366357
"outputs": [],
367358
"source": [
368-
"rag_file = rag.upload_file(\n",
359+
"rag_file = client.rag.upload_file(\n",
369360
" corpus_name=rag_corpus.name,\n",
370361
" path=\"test.md\",\n",
371362
" display_name=\"test.md\",\n",
372-
" description=\"my test file\",\n",
373363
")"
374364
]
375365
},
@@ -395,17 +385,21 @@
395385
"outputs": [],
396386
"source": [
397387
"INPUT_GCS_BUCKET = (\n",
398-
" \"gs://cloud-samples-data/gen-app-builder/search/alphabet-investor-pdfs/\"\n",
388+
" \"gs://cloud-samples-data/gen-app-builder/search/alphabet-investor-pdfs/*\"\n",
399389
")\n",
400390
"\n",
401-
"response = rag.import_files(\n",
402-
" corpus_name=rag_corpus.name,\n",
403-
" paths=[INPUT_GCS_BUCKET],\n",
404-
" # Optional\n",
405-
" transformation_config=rag.TransformationConfig(\n",
406-
" chunking_config=rag.ChunkingConfig(chunk_size=1024, chunk_overlap=100)\n",
391+
"response = client.rag.import_files(\n",
392+
" name=rag_corpus.name,\n",
393+
" import_config=types.ImportRagFilesConfig(\n",
394+
" gcs_source=types.GcsSource(uris=[INPUT_GCS_BUCKET]),\n",
395+
" rag_file_transformation_config=types.RagFileTransformationConfig(\n",
396+
" rag_file_chunking_config=types.RagFileChunkingConfig(\n",
397+
" chunk_size=1024,\n",
398+
" chunk_overlap=100,\n",
399+
" )\n",
400+
" ), # optional\n",
401+
" max_embedding_requests_per_min=900, # optional\n",
407402
" ),\n",
408-
" max_embedding_requests_per_min=900, # Optional\n",
409403
")"
410404
]
411405
},
@@ -433,12 +427,23 @@
433427
},
434428
"outputs": [],
435429
"source": [
436-
"response = rag.import_files(\n",
437-
" corpus_name=rag_corpus.name,\n",
438-
" paths=[\"https://drive.google.com/drive/folders/{folder_id}\"],\n",
439-
" # Optional\n",
440-
" transformation_config=rag.TransformationConfig(\n",
441-
" chunking_config=rag.ChunkingConfig(chunk_size=512, chunk_overlap=50)\n",
430+
"response = client.rag.import_files(\n",
431+
" name=rag_corpus.name,\n",
432+
" import_config=types.ImportRagFilesConfig(\n",
433+
" google_drive_source=types.GoogleDriveSource(\n",
434+
" resource_ids=[\n",
435+
" types.GoogleDriveSourceResourceId(\n",
436+
" resource_id=\"https://drive.google.com/drive/folders/{folder_id}\",\n",
437+
" resource_type=types.ResourceType.RESOURCE_TYPE_FOLDER,\n",
438+
" )\n",
439+
" ]\n",
440+
" ),\n",
441+
" rag_file_transformation_config=types.RagFileTransformationConfig(\n",
442+
" rag_file_chunking_config=types.RagFileChunkingConfig(\n",
443+
" chunk_size=512,\n",
444+
" chunk_overlap=50,\n",
445+
" )\n",
446+
" ), # optional\n",
442447
" ),\n",
443448
")"
444449
]
@@ -461,21 +466,19 @@
461466
"outputs": [],
462467
"source": [
463468
"# Direct context retrieval\n",
464-
"response = rag.retrieval_query(\n",
465-
" rag_resources=[\n",
466-
" rag.RagResource(\n",
467-
" rag_corpus=rag_corpus.name,\n",
468-
" # Optional: supply IDs from `rag.list_files()`.\n",
469-
" # rag_file_ids=[\"rag-file-1\", \"rag-file-2\", ...],\n",
470-
" )\n",
471-
" ],\n",
472-
" rag_retrieval_config=rag.RagRetrievalConfig(\n",
473-
" top_k=10, # Optional\n",
474-
" filter=rag.Filter(\n",
475-
" vector_distance_threshold=0.5, # Optional\n",
469+
"response = client.rag.retrieve_contexts(\n",
470+
" vertex_rag_store=genai_types.VertexRagStore(\n",
471+
" rag_resources=[\n",
472+
" genai_types.VertexRagStoreRagResource(rag_corpus=rag_corpus.name)\n",
473+
" ]\n",
474+
" ),\n",
475+
" query=types.RagQuery(\n",
476+
" text=\"What is RAG and why it is helpful?\",\n",
477+
" rag_retrieval_config=types.RagRetrievalConfig(\n",
478+
" top_k=10,\n",
479+
" filter=types.RagRetrievalConfigFilter(vector_distance_threshold=0.5),\n",
476480
" ),\n",
477481
" ),\n",
478-
" text=\"What is RAG and why it is helpful?\",\n",
479482
")\n",
480483
"print(response)\n",
481484
"\n",
@@ -501,9 +504,9 @@
501504
"outputs": [],
502505
"source": [
503506
"# Create a tool for the RAG Corpus\n",
504-
"rag_retrieval_tool = Tool(\n",
505-
" retrieval=Retrieval(\n",
506-
" vertex_rag_store=VertexRagStore(\n",
507+
"rag_retrieval_tool = genai_types.Tool(\n",
508+
" retrieval=genai_types.Retrieval(\n",
509+
" vertex_rag_store=genai_types.VertexRagStore(\n",
507510
" rag_corpora=[rag_corpus.name],\n",
508511
" similarity_top_k=10,\n",
509512
" vector_distance_threshold=0.5,\n",
@@ -536,70 +539,28 @@
536539
"cell_type": "code",
537540
"execution_count": null,
538541
"metadata": {
539-
"id": "124b36be8d5b"
542+
"id": "1f02122723b6"
540543
},
541544
"outputs": [],
542545
"source": [
543-
"response = client.models.generate_content(\n",
544-
" model=MODEL_ID,\n",
545-
" contents=\"What is RAG?\",\n",
546-
" config=GenerateContentConfig(tools=[rag_retrieval_tool]),\n",
547-
")\n",
548-
"\n",
549-
"display(Markdown(response.text))"
550-
]
551-
},
552-
{
553-
"cell_type": "markdown",
554-
"metadata": {
555-
"id": "0268fe43d41c"
556-
},
557-
"source": [
558-
"### Generate Content with Llama3 using RAG Retrieval Tool"
546+
"# Create a Gen AI SDK client\n",
547+
"genai_client = genai.Client(project=PROJECT_ID, location=\"your-location\")"
559548
]
560549
},
561550
{
562551
"cell_type": "code",
563552
"execution_count": null,
564553
"metadata": {
565-
"id": "f6e67ee7968c"
554+
"id": "124b36be8d5b"
566555
},
567556
"outputs": [],
568557
"source": [
569-
"from vertexai import generative_models\n",
570-
"\n",
571-
"# Load tool into Llama model\n",
572-
"rag_retrieval_tool = generative_models.Tool.from_retrieval(\n",
573-
" retrieval=rag.Retrieval(\n",
574-
" source=rag.VertexRagStore(\n",
575-
" rag_resources=[rag.RagResource(rag_corpus=rag_corpus.name)],\n",
576-
" rag_retrieval_config=rag.RagRetrievalConfig(\n",
577-
" top_k=10, # Optional\n",
578-
" filter=rag.Filter(\n",
579-
" vector_distance_threshold=0.5, # Optional\n",
580-
" ),\n",
581-
" ),\n",
582-
" ),\n",
583-
" )\n",
558+
"response = genai_client.models.generate_content(\n",
559+
" model=MODEL_ID,\n",
560+
" contents=\"What is RAG?\",\n",
561+
" config=genai_types.GenerateContentConfig(tools=[rag_retrieval_tool]),\n",
584562
")\n",
585563
"\n",
586-
"llama_model = generative_models.GenerativeModel(\n",
587-
" # your self-deployed endpoint for Llama3\n",
588-
" \"projects/{project}/locations/{location}/endpoints/{endpoint_resource_id}\",\n",
589-
" tools=[rag_retrieval_tool],\n",
590-
")"
591-
]
592-
},
593-
{
594-
"cell_type": "code",
595-
"execution_count": null,
596-
"metadata": {
597-
"id": "c6d710b6dece"
598-
},
599-
"outputs": [],
600-
"source": [
601-
"response = llama_model.generate_content(\"What is RAG?\")\n",
602-
"\n",
603564
"display(Markdown(response.text))"
604565
]
605566
}

0 commit comments

Comments
 (0)