more tweaks to TEST_GAPIC_RAG_CORPUS

Author: chalmerlowe

tests/unit/vertex_rag/test_rag_constants_preview.py

@@ -17,14 +17,15 @@
 
 
 from google.cloud import aiplatform
-from google.cloud.aiplatform_v1beta1 import (
+from google.cloud.aiplatform_v1beta1.types import (
     GoogleDriveSource,
     ImportRagFilesConfig,
     ImportRagFilesRequest,
     ImportRagFilesResponse,
     JiraSource as GapicJiraSource,
     RagContexts,
     RagCorpus as GapicRagCorpus,
+    RagEmbeddingModelConfig as GapicRagEmbeddingModelConfig,
     RagEngineConfig as GapicRagEngineConfig,
     RagFileChunkingConfig,
     RagFileParsingConfig,
@@ -36,9 +37,9 @@
     SharePointSources as GapicSharePointSources,
     SlackSource as GapicSlackSource,
     VertexAiSearchConfig as GapicVertexAiSearchConfig,
+    api_auth,
+    EncryptionSpec,
 )
-from google.cloud.aiplatform_v1beta1.types import api_auth
-from google.cloud.aiplatform_v1beta1.types import EncryptionSpec
 from vertexai.preview.rag import (
     ANN,
     Basic,
@@ -137,12 +138,10 @@
     name=TEST_RAG_CORPUS_RESOURCE_NAME,
     display_name=TEST_CORPUS_DISPLAY_NAME,
     description=TEST_CORPUS_DISCRIPTION,
-    vector_db_config=GapicRagVectorDbConfig(
-        rag_embedding_model_config=GapicRagVectorDbConfig.RagEmbeddingModelConfig(
-            vertex_prediction_endpoint=GapicRagVectorDbConfig.RagEmbeddingModelConfig.VertexPredictionEndpoint(
-                endpoint="projects/{}/locations/{}/publishers/google/models/textembedding-gecko".format(
-                    TEST_PROJECT, TEST_REGION
-                )
+    rag_embedding_model_config=GapicRagEmbeddingModelConfig(
+        vertex_prediction_endpoint=GapicRagEmbeddingModelConfig.VertexPredictionEndpoint(
+            endpoint="projects/{}/locations/{}/publishers/google/models/textembedding-gecko".format(
+                TEST_PROJECT, TEST_REGION
             )
         )
     ),