feat: Update preview create_corpus to accept encryption_spec in rag_data.py

PiperOrigin-RevId: 761091381

Author: vertex-sdk-bot

tests/unit/vertex_rag/test_rag_constants_preview.py

@@ -38,6 +38,7 @@
  VertexAiSearchConfig as GapicVertexAiSearchConfig,
 )
 from google.cloud.aiplatform_v1beta1.types import api_auth
+from google.cloud.aiplatform_v1beta1.types import EncryptionSpec
 from vertexai.preview.rag import (
  ANN,
  Basic,
@@ -91,6 +92,9 @@
 TEST_WEAVIATE_API_KEY_SECRET_VERSION = (
  "projects/test-project/secrets/test-secret/versions/1"
 )
+TEST_ENCRYPTION_SPEC = EncryptionSpec(
+ kms_key_name="projects/test-project/locations/us-central1/keyRings/test-key-ring/cryptoKeys/test-key"
+)
 TEST_WEAVIATE_CONFIG = Weaviate(
  weaviate_http_endpoint=TEST_WEAVIATE_HTTP_ENDPOINT,
  collection_name=TEST_WEAVIATE_COLLECTION_NAME,
@@ -133,6 +137,14 @@
  TEST_PROJECT, TEST_REGION
  )
 )
+TEST_GAPIC_CMEK_RAG_CORPUS = GapicRagCorpus(
+ name=TEST_RAG_CORPUS_RESOURCE_NAME,
+ display_name=TEST_CORPUS_DISPLAY_NAME,
+ description=TEST_CORPUS_DISCRIPTION,
+ encryption_spec=EncryptionSpec(
+ kms_key_name="projects/test-project/locations/us-central1/keyRings/test-key-ring/cryptoKeys/test-key"
+ ),
+)
 TEST_GAPIC_RAG_CORPUS_WEAVIATE = GapicRagCorpus(
  name=TEST_RAG_CORPUS_RESOURCE_NAME,
  display_name=TEST_CORPUS_DISPLAY_NAME,
@@ -236,6 +248,14 @@
  description=TEST_CORPUS_DISCRIPTION,
  embedding_model_config=TEST_EMBEDDING_MODEL_CONFIG,
 )
+TEST_CMEK_RAG_CORPUS = RagCorpus(
+ name=TEST_RAG_CORPUS_RESOURCE_NAME,
+ display_name=TEST_CORPUS_DISPLAY_NAME,
+ description=TEST_CORPUS_DISCRIPTION,
+ encryption_spec=EncryptionSpec(
+ kms_key_name="projects/test-project/locations/us-central1/keyRings/test-key-ring/cryptoKeys/test-key"
+ ),
+)
 TEST_RAG_CORPUS_WEAVIATE = RagCorpus(
  name=TEST_RAG_CORPUS_RESOURCE_NAME,
  display_name=TEST_CORPUS_DISPLAY_NAME,

tests/unit/vertex_rag/test_rag_data_preview.py

@@ -53,6 +53,21 @@ def create_rag_corpus_mock():
  yield create_rag_corpus_mock
 
 
+@pytest.fixture
+def create_rag_corpus_mock_cmek():
+ with mock.patch.object(
+ VertexRagDataServiceClient,
+ "create_rag_corpus",
+ ) as create_rag_corpus_mock_cmek:
+ create_rag_corpus_lro_mock = mock.Mock(ga_operation.Operation)
+ create_rag_corpus_lro_mock.done.return_value = True
+ create_rag_corpus_lro_mock.result.return_value = (
+ test_rag_constants_preview.TEST_GAPIC_CMEK_RAG_CORPUS
+ )
+ create_rag_corpus_mock_cmek.return_value = create_rag_corpus_lro_mock
+ yield create_rag_corpus_mock_cmek
+
+
 @pytest.fixture
 def create_rag_corpus_mock_backend():
  with mock.patch.object(
@@ -653,6 +668,15 @@ def test_create_corpus_backend_success(self):
 
  rag_corpus_eq(rag_corpus, test_rag_constants_preview.TEST_RAG_CORPUS_BACKEND)
 
+ @pytest.mark.usefixtures("create_rag_corpus_mock_cmek")
+ def test_create_corpus_cmek_success(self):
+ rag_corpus = rag.create_corpus(
+ display_name=test_rag_constants_preview.TEST_CORPUS_DISPLAY_NAME,
+ encryption_spec=test_rag_constants_preview.TEST_ENCRYPTION_SPEC,
+ )
+
+ rag_corpus_eq(rag_corpus, test_rag_constants_preview.TEST_CMEK_RAG_CORPUS)
+
  @pytest.mark.usefixtures("create_rag_corpus_mock_weaviate")
  def test_create_corpus_weaviate_success(self):
  rag_corpus = rag.create_corpus(

vertexai/preview/rag/rag_data.py

@@ -41,6 +41,7 @@
  ListRagCorporaPager,
  ListRagFilesPager,
 )
+from google.cloud.aiplatform_v1beta1.types import EncryptionSpec
 from vertexai.preview.rag.utils import (
  _gapic_utils,
 )
@@ -74,6 +75,7 @@ def create_corpus(
  ] = None,
  vertex_ai_search_config: Optional[VertexAiSearchConfig] = None,
  backend_config: Optional[RagVectorDbConfig] = None,
+ encryption_spec: Optional[EncryptionSpec] = None,
 ) -> RagCorpus:
  """Creates a new RagCorpus resource.
 
@@ -91,8 +93,8 @@ def create_corpus(
 
  Args:
  display_name: If not provided, SDK will create one. The display name of
- the RagCorpus. The name can be up to 128 characters long and can
- consist of any UTF-8 characters.
+ the RagCorpus. The name can be up to 128 characters long and can consist
+ of any UTF-8 characters.
  description: The description of the RagCorpus.
  embedding_model_config: The embedding model config.
  Note: Deprecated. Use backend_config instead.
@@ -104,6 +106,8 @@ def create_corpus(
  vertex_ai_search_config is specified.
  backend_config: The backend config of the RagCorpus. It can specify a
  Vector DB and/or the embedding model config.
+ encryption_spec: The encryption spec of the RagCorpus.
+
  Returns:
  RagCorpus.
  Raises:
@@ -156,6 +160,12 @@ def create_corpus(
  rag_corpus=rag_corpus,
  )
 
+ if encryption_spec:
+ _gapic_utils.set_encryption_spec(
+ encryption_spec=encryption_spec,
+ rag_corpus=rag_corpus,
+ )
+
  request = CreateRagCorpusRequest(
  parent=parent,
  rag_corpus=rag_corpus,

vertexai/preview/rag/utils/_gapic_utils.py

@@ -41,6 +41,7 @@
  VertexAiSearchConfig as GapicVertexAiSearchConfig,
 )
 from google.cloud.aiplatform_v1beta1.types import api_auth
+from google.cloud.aiplatform_v1beta1.types import EncryptionSpec
 from vertexai.preview.rag.utils.resources import (
  ANN,
  EmbeddingModelConfig,
@@ -327,6 +328,7 @@ def convert_gapic_to_rag_corpus(gapic_rag_corpus: GapicRagCorpus) -> RagCorpus:
  backend_config=convert_gapic_to_backend_config(
  gapic_rag_corpus.vector_db_config
  ),
+ encryption_spec=gapic_rag_corpus.encryption_spec,
  )
  return rag_corpus
 
@@ -348,6 +350,7 @@ def convert_gapic_to_rag_corpus_no_embedding_model_config(
  backend_config=convert_gapic_to_backend_config(
  rag_vector_db_config_no_embedding_model_config
  ),
+ encryption_spec=gapic_rag_corpus.encryption_spec,
  )
  return rag_corpus
 
@@ -728,6 +731,28 @@ def set_embedding_model_config(
  )
 
 
+def set_encryption_spec(
+ encryption_spec: EncryptionSpec,
+ rag_corpus: GapicRagCorpus,
+) -> None:
+ """Sets the encryption spec for the rag corpus."""
+ # Raises value error if encryption_spec.kms_key_name is None or empty,
+ if encryption_spec.kms_key_name is None or not encryption_spec.kms_key_name:
+ raise ValueError("kms_key_name must be set if encryption_spec is set.")
+
+ # Raises value error if encryption_spec.kms_key_name is not a valid KMS key name.
+ if not re.match(
+ r"^projects/(?P<project>.+?)/locations/(?P<location>.+?)/keyRings/(?P<key_ring>.+?)/cryptoKeys/(?P<crypto_key>.+?)$",
+ encryption_spec.kms_key_name,
+ ):
+ raise ValueError(
+ "kms_key_name must be of the format "
+ "`projects/{project}/locations/{location}/keyRings/{key_ring}/cryptoKeys/{crypto_key}`"
+ )
+
+ rag_corpus.encryption_spec = encryption_spec
+
+
 def set_vector_db(
  vector_db: Union[
  Weaviate, VertexFeatureStore, VertexVectorSearch, Pinecone, RagManagedDb, None

vertexai/preview/rag/utils/resources.py

@@ -18,6 +18,8 @@
 import dataclasses
 from typing import List, Optional, Sequence, Union
 
+from google.cloud.aiplatform_v1beta1.types import EncryptionSpec
+
 from google.protobuf import timestamp_pb2
 
 DEPRECATION_DATE = "June 2025"
@@ -270,6 +272,7 @@ class RagCorpus:
  vertex_ai_search_config: The Vertex AI Search config of the RagCorpus.
  backend_config: The backend config of the RagCorpus. It can specify a
  Vector DB and/or the embedding model config.
+ encryption_spec: The encryption spec of the RagCorpus. Immutable.
  """
 
  name: Optional[str] = None
@@ -281,6 +284,7 @@ class RagCorpus:
  ] = None
  vertex_ai_search_config: Optional[VertexAiSearchConfig] = None
  backend_config: Optional[RagVectorDbConfig] = None
+ encryption_spec: Optional[EncryptionSpec] = None
 
 
 @dataclasses.dataclass