feat: Refactor rag_store and rag_retrieval to use v1 protos

PiperOrigin-RevId: 700410303

google/cloud/aiplatform/compat/services/__init__.py

@@ -179,6 +179,15 @@
 from google.cloud.aiplatform_v1.services.vizier_service import (
     client as vizier_service_client_v1,
 )
+from google.cloud.aiplatform_v1.services.vertex_rag_service import (
+    client as vertex_rag_service_client_v1,
+)
+from google.cloud.aiplatform_v1.services.vertex_rag_data_service import (
+    client as vertex_rag_data_service_client_v1,
+)
+from google.cloud.aiplatform_v1.services.vertex_rag_data_service import (
+    async_client as vertex_rag_data_service_async_client_v1,
+)
 
 __all__ = (
     # v1
@@ -204,6 +213,9 @@
     specialist_pool_service_client_v1,
     tensorboard_service_client_v1,
     vizier_service_client_v1,
+    vertex_rag_data_service_async_client_v1,
+    vertex_rag_data_service_client_v1,
+    vertex_rag_service_client_v1,
     # v1beta1
     dataset_service_client_v1beta1,
     deployment_resource_pool_service_client_v1beta1,

google/cloud/aiplatform/utils/__init__.py

@@ -90,6 +90,9 @@
     tensorboard_service_client_v1,
     vizier_service_client_v1,
     persistent_resource_service_client_v1,
+    vertex_rag_data_service_async_client_v1,
+    vertex_rag_data_service_client_v1,
+    vertex_rag_service_client_v1,
 )
 
 from google.cloud.aiplatform.compat.types import (
@@ -138,6 +141,7 @@
     schedule_service_client_v1.ScheduleServiceClient,
     tensorboard_service_client_v1.TensorboardServiceClient,
     vizier_service_client_v1.VizierServiceClient,
+    vertex_rag_service_client_v1.VertexRagServiceClient,
 )
 
 
@@ -967,8 +971,9 @@ class ReasoningEngineExecutionClientWithOverride(ClientWithOverride):
 
 class VertexRagDataClientWithOverride(ClientWithOverride):
     _is_temporary = True
-    _default_version = compat.V1BETA1
+    _default_version = compat.DEFAULT_VERSION
     _version_map = (
+        (compat.V1, vertex_rag_data_service_client_v1.VertexRagDataServiceClient),
         (
             compat.V1BETA1,
             vertex_rag_data_service_client_v1beta1.VertexRagDataServiceClient,
@@ -978,8 +983,12 @@ class VertexRagDataClientWithOverride(ClientWithOverride):
 
 class VertexRagDataAsyncClientWithOverride(ClientWithOverride):
     _is_temporary = True
-    _default_version = compat.V1BETA1
+    _default_version = compat.DEFAULT_VERSION
     _version_map = (
+        (
+            compat.V1,
+            vertex_rag_data_service_async_client_v1.VertexRagDataServiceAsyncClient,
+        ),
         (
             compat.V1BETA1,
             vertex_rag_data_service_async_client_v1beta1.VertexRagDataServiceAsyncClient,
@@ -989,12 +998,10 @@ class VertexRagDataAsyncClientWithOverride(ClientWithOverride):
 
 class VertexRagClientWithOverride(ClientWithOverride):
     _is_temporary = True
-    _default_version = compat.V1BETA1
+    _default_version = compat.DEFAULT_VERSION
     _version_map = (
-        (
-            compat.V1BETA1,
-            vertex_rag_service_client_v1beta1.VertexRagServiceClient,
-        ),
+        (compat.V1, vertex_rag_service_client_v1.VertexRagServiceClient),
+        (compat.V1BETA1, vertex_rag_service_client_v1beta1.VertexRagServiceClient),
     )
 
 

tests/unit/vertex_rag/test_rag_constants.py

@@ -18,12 +18,14 @@
 
 from google.cloud import aiplatform
 
-from vertexai.preview.rag import (
+from vertexai.rag import (
     EmbeddingModelConfig,
+    Filter,
     Pinecone,
     RagCorpus,
     RagFile,
     RagResource,
+    RagRetrievalConfig,
     SharePointSource,
     SharePointSources,
     SlackChannelsSource,
@@ -34,6 +36,7 @@
     VertexVectorSearch,
     VertexFeatureStore,
 )
+
 from google.cloud.aiplatform_v1beta1 import (
     GoogleDriveSource,
     RagFileChunkingConfig,
@@ -46,9 +49,11 @@
     RagFile as GapicRagFile,
     SharePointSources as GapicSharePointSources,
     SlackSource as GapicSlackSource,
+    RagVectorDbConfig,
+)
+from google.cloud.aiplatform_v1 import (
     RagContexts,
     RetrieveContextsResponse,
-    RagVectorDbConfig,
 )
 from google.cloud.aiplatform_v1beta1.types import api_auth
 from google.protobuf import timestamp_pb2
@@ -529,3 +534,7 @@
     rag_corpus="213lkj-1/23jkl/",
     rag_file_ids=[TEST_RAG_FILE_ID],
 )
+TEST_RAG_RETRIEVAL_CONFIG = RagRetrievalConfig(
+    top_k=2,
+    filter=Filter(vector_distance_threshold=0.5),
+)