docs(generative_ai): Update RAG Engine Samples to V1 (#13287)

holtskinner · web-flow · commit 46785e2b83d4 · 2025-04-04T13:53:19.000+11:00
* docs(generative_ai): Update RAG Engine Samples to V1

* Formatting

* Undo test changes
diff --git a/generative_ai/rag/create_corpus_example.py b/generative_ai/rag/create_corpus_example.py
@@ -26,7 +26,7 @@ def create_corpus(
 ) -> RagCorpus:
     # [START generativeaionvertexai_rag_create_corpus]
 
-    from vertexai.preview import rag
+    from vertexai import rag
     import vertexai
 
     # TODO(developer): Update and un-comment below lines
@@ -37,15 +37,19 @@ def create_corpus(
     # Initialize Vertex AI API once per session
     vertexai.init(project=PROJECT_ID, location="us-central1")
 
-    # Configure embedding model (Optional)
-    embedding_model_config = rag.EmbeddingModelConfig(
-        publisher_model="publishers/google/models/text-embedding-004"
+    # 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"
+            )
+        )
     )
 
     corpus = rag.create_corpus(
         display_name=display_name,
         description=description,
-        embedding_model_config=embedding_model_config,
+        backend_config=backend_config,
     )
     print(corpus)
     # Example response:
diff --git a/generative_ai/rag/create_corpus_pinecone_example.py b/generative_ai/rag/create_corpus_pinecone_example.py
@@ -28,7 +28,7 @@ def create_corpus_pinecone(
 ) -> RagCorpus:
     # [START generativeaionvertexai_rag_create_corpus_pinecone]
 
-    from vertexai.preview import rag
+    from vertexai import rag
     import vertexai
 
     # TODO(developer): Update and un-comment below lines
@@ -42,8 +42,10 @@ def create_corpus_pinecone(
     vertexai.init(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 = rag.RagEmbeddingModelConfig(
+        vertex_prediction_endpoint=rag.VertexPredictionEndpoint(
+            publisher_model="publishers/google/models/text-embedding-005"
+        )
     )
 
     # Configure Vector DB
@@ -55,8 +57,10 @@ def create_corpus_pinecone(
     corpus = rag.create_corpus(
         display_name=display_name,
         description=description,
-        embedding_model_config=embedding_model_config,
-        vector_db=vector_db,
+        backend_config=rag.RagVectorDbConfig(
+            rag_embedding_model_config=embedding_model_config,
+            vector_db=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
@@ -28,7 +28,7 @@ def create_corpus_vector_search(
 ) -> RagCorpus:
     # [START generativeaionvertexai_rag_create_corpus_vector_search]
 
-    from vertexai.preview import rag
+    from vertexai import rag
     import vertexai
 
     # TODO(developer): Update and un-comment below lines
@@ -42,8 +42,10 @@ def create_corpus_vector_search(
     vertexai.init(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 = rag.RagEmbeddingModelConfig(
+        vertex_prediction_endpoint=rag.VertexPredictionEndpoint(
+            publisher_model="publishers/google/models/text-embedding-005"
+        )
     )
 
     # Configure Vector DB
@@ -54,8 +56,10 @@ def create_corpus_vector_search(
     corpus = rag.create_corpus(
         display_name=display_name,
         description=description,
-        embedding_model_config=embedding_model_config,
-        vector_db=vector_db,
+        backend_config=rag.RagVectorDbConfig(
+            rag_embedding_model_config=embedding_model_config,
+            vector_db=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
@@ -15,7 +15,7 @@
 
 from typing import Optional
 
-from vertexai.preview.rag import RagCorpus
+from vertexai import rag
 
 PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
 
@@ -24,10 +24,10 @@ def create_corpus_vertex_ai_search(
     vertex_ai_search_engine_name: str,
     display_name: Optional[str] = None,
     description: Optional[str] = None,
-) -> RagCorpus:
+) -> rag.RagCorpus:
     # [START generativeaionvertexai_rag_create_corpus_vertex_ai_search]
 
-    from vertexai.preview import rag
+    from vertexai import rag
     import vertexai
 
     # TODO(developer): Update and un-comment below lines
diff --git a/generative_ai/rag/delete_corpus_example.py b/generative_ai/rag/delete_corpus_example.py
@@ -20,7 +20,7 @@
 def delete_corpus(corpus_name: str) -> None:
     # [START generativeaionvertexai_rag_delete_corpus]
 
-    from vertexai.preview import rag
+    from vertexai import rag
     import vertexai
 
     # TODO(developer): Update and un-comment below lines
diff --git a/generative_ai/rag/delete_file_example.py b/generative_ai/rag/delete_file_example.py
@@ -20,7 +20,7 @@
 def delete_file(file_name: str) -> None:
     # [START generativeaionvertexai_rag_delete_file]
 
-    from vertexai.preview import rag
+    from vertexai import rag
     import vertexai
 
     # TODO(developer): Update and un-comment below lines
diff --git a/generative_ai/rag/generate_content_example.py b/generative_ai/rag/generate_content_example.py
@@ -24,8 +24,8 @@ def generate_content_with_rag(
 ) -> GenerationResponse:
     # [START generativeaionvertexai_rag_generate_content]
 
-    from vertexai.preview import rag
-    from vertexai.preview.generative_models import GenerativeModel, Tool
+    from vertexai import rag
+    from vertexai.generative_models import GenerativeModel, Tool
     import vertexai
 
     # TODO(developer): Update and un-comment below lines
@@ -45,14 +45,16 @@ def generate_content_with_rag(
                         # rag_file_ids=["rag-file-1", "rag-file-2", ...],
                     )
                 ],
-                similarity_top_k=3,  # Optional
-                vector_distance_threshold=0.5,  # Optional
+                rag_retrieval_config=rag.RagRetrievalConfig(
+                    top_k=10,
+                    filter=rag.utils.resources.Filter(vector_distance_threshold=0.5),
+                ),
             ),
         )
     )
 
     rag_model = GenerativeModel(
-        model_name="gemini-1.5-flash-001", tools=[rag_retrieval_tool]
+        model_name="gemini-2.0-flash-001", tools=[rag_retrieval_tool]
     )
     response = rag_model.generate_content("Why is the sky blue?")
     print(response.text)
diff --git a/generative_ai/rag/get_corpus_example.py b/generative_ai/rag/get_corpus_example.py
@@ -22,7 +22,7 @@
 def get_corpus(corpus_name: str) -> RagCorpus:
     # [START generativeaionvertexai_rag_get_corpus]
 
-    from vertexai.preview import rag
+    from vertexai import rag
     import vertexai
 
     # TODO(developer): Update and un-comment below lines
diff --git a/generative_ai/rag/get_file_example.py b/generative_ai/rag/get_file_example.py
@@ -14,15 +14,15 @@
 
 import os
 
-from google.cloud.aiplatform_v1beta1 import RagFile
+from google.cloud.aiplatform_v1 import RagFile
 
 PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
 
 
 def get_file(file_name: str) -> RagFile:
     # [START generativeaionvertexai_rag_get_file]
 
-    from vertexai.preview import rag
+    from vertexai import rag
     import vertexai
 
     # TODO(developer): Update and un-comment below lines
diff --git a/generative_ai/rag/import_files_async_example.py b/generative_ai/rag/import_files_async_example.py
@@ -16,7 +16,7 @@
 
 from typing import List
 
-from google.cloud.aiplatform_v1beta1 import ImportRagFilesResponse
+from google.cloud.aiplatform_v1 import ImportRagFilesResponse
 
 PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
 
@@ -27,7 +27,7 @@ async def import_files_async(
 ) -> ImportRagFilesResponse:
     # [START generativeaionvertexai_rag_import_files_async]
 
-    from vertexai.preview import rag
+    from vertexai import rag
     import vertexai
 
     # TODO(developer): Update and un-comment below lines
@@ -40,11 +40,12 @@ async def import_files_async(
     # Initialize Vertex AI API once per session
     vertexai.init(project=PROJECT_ID, location="us-central1")
 
-    response = await rag.import_files_async(
+    response = await rag.import_files(
         corpus_name=corpus_name,
         paths=paths,
-        chunk_size=512,  # Optional
-        chunk_overlap=100,  # Optional
+        transformation_config=rag.TransformationConfig(
+            rag.ChunkingConfig(chunk_size=512, chunk_overlap=100)
+        ),
         max_embedding_requests_per_min=900,  # Optional
     )
 
diff --git a/generative_ai/rag/import_files_example.py b/generative_ai/rag/import_files_example.py
@@ -15,7 +15,7 @@
 import os
 from typing import List
 
-from google.cloud.aiplatform_v1beta1 import ImportRagFilesResponse
+from google.cloud.aiplatform_v1 import ImportRagFilesResponse
 
 PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
 
@@ -26,7 +26,7 @@ def import_files(
 ) -> ImportRagFilesResponse:
     # [START generativeaionvertexai_rag_import_files]
 
-    from vertexai.preview import rag
+    from vertexai import rag
     import vertexai
 
     # TODO(developer): Update and un-comment below lines
@@ -40,8 +40,9 @@ def import_files(
     response = rag.import_files(
         corpus_name=corpus_name,
         paths=paths,
-        chunk_size=512,  # Optional
-        chunk_overlap=100,  # Optional
+        transformation_config=rag.TransformationConfig(
+            rag.ChunkingConfig(chunk_size=512, chunk_overlap=100)
+        ),
         max_embedding_requests_per_min=900,  # Optional
     )
     print(f"Imported {response.imported_rag_files_count} files.")
diff --git a/generative_ai/rag/list_corpora_example.py b/generative_ai/rag/list_corpora_example.py
@@ -24,7 +24,7 @@
 def list_corpora() -> ListRagCorporaPager:
     # [START generativeaionvertexai_rag_list_corpora]
 
-    from vertexai.preview import rag
+    from vertexai import rag
     import vertexai
 
     # TODO(developer): Update and un-comment below lines
diff --git a/generative_ai/rag/list_files_example.py b/generative_ai/rag/list_files_example.py
@@ -14,7 +14,7 @@
 
 import os
 
-from google.cloud.aiplatform_v1beta1.services.vertex_rag_data_service.pagers import (
+from google.cloud.aiplatform_v1.services.vertex_rag_data_service.pagers import (
     ListRagFilesPager,
 )
 
@@ -24,7 +24,7 @@
 def list_files(corpus_name: str) -> ListRagFilesPager:
     # [START generativeaionvertexai_rag_list_files]
 
-    from vertexai.preview import rag
+    from vertexai import rag
     import vertexai
 
     # TODO(developer): Update and un-comment below lines
diff --git a/generative_ai/rag/quickstart_example.py b/generative_ai/rag/quickstart_example.py
@@ -16,19 +16,19 @@
 
 from typing import List, Tuple
 
+from vertexai import rag
 from vertexai.generative_models import GenerationResponse
-from vertexai.preview.rag.utils.resources import RagCorpus
 
 PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
 
 
 def quickstart(
     display_name: str,
     paths: List[str],
-) -> Tuple[RagCorpus, GenerationResponse]:
+) -> Tuple[rag.RagCorpus, GenerationResponse]:
     # [START generativeaionvertexai_rag_quickstart]
-    from vertexai.preview import rag
-    from vertexai.preview.generative_models import GenerativeModel, Tool
+    from vertexai import rag
+    from vertexai.generative_models import GenerativeModel, Tool
     import vertexai
 
     # Create a RAG Corpus, Import Files, and Generate a response
@@ -42,26 +42,39 @@ def quickstart(
     vertexai.init(project=PROJECT_ID, location="us-central1")
 
     # Create RagCorpus
-    # Configure embedding model, for example "text-embedding-004".
-    embedding_model_config = rag.EmbeddingModelConfig(
-        publisher_model="publishers/google/models/text-embedding-004"
+    # 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,
-        embedding_model_config=embedding_model_config,
+        backend_config=rag.RagVectorDbConfig(
+            rag_embedding_model_config=embedding_model_config
+        ),
     )
 
     # Import Files to the RagCorpus
     rag.import_files(
         rag_corpus.name,
         paths,
-        chunk_size=512,  # Optional
-        chunk_overlap=100,  # Optional
-        max_embedding_requests_per_min=900,  # Optional
+        # 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(
@@ -71,8 +84,7 @@ def quickstart(
             )
         ],
         text="What is RAG and why it is helpful?",
-        similarity_top_k=10,  # Optional
-        vector_distance_threshold=0.5,  # Optional
+        rag_retrieval_config=rag_retrieval_config,
     )
     print(response)
 
@@ -88,14 +100,14 @@ def quickstart(
                         # rag_file_ids=["rag-file-1", "rag-file-2", ...],
                     )
                 ],
-                similarity_top_k=3,  # Optional
-                vector_distance_threshold=0.5,  # Optional
+                rag_retrieval_config=rag_retrieval_config,
             ),
         )
     )
-    # Create a gemini-pro model instance
+
+    # Create a Gemini model instance
     rag_model = GenerativeModel(
-        model_name="gemini-1.5-flash-001", tools=[rag_retrieval_tool]
+        model_name="gemini-2.0-flash-001", tools=[rag_retrieval_tool]
     )
 
     # Generate response
diff --git a/generative_ai/rag/requirements.txt b/generative_ai/rag/requirements.txt
@@ -1 +1 @@
-google-cloud-aiplatform[all]==1.74.0
+google-cloud-aiplatform==1.87.0
diff --git a/generative_ai/rag/retrieval_query_example.py b/generative_ai/rag/retrieval_query_example.py
diff --git a/generative_ai/rag/upload_file_example.py b/generative_ai/rag/upload_file_example.py
