8000 feat: Update v1 `create_corpus` to accept `encryption_spec` in `rag_d… · googleapis/python-aiplatform@7a317a5 · GitHub
[go: up one dir, main page]
Skip to content

Commit 7a317a5

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
feat: Update v1 create_corpus to accept encryption_spec in rag_data.py
PiperOrigin-RevId: 772502746
1 parent 0059c01 commit 7a317a5

File tree

5 files changed

+81
-0
lines changed

5 files changed

+81
-0
lines changed

‎tests/unit/vertex_rag/test_rag_constants.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,7 @@
6262
VertexAiSearchConfig as GapicVertexAiSearchConfig,
6363
)
6464
from google.cloud.aiplatform_v1.types import api_auth
65+
from google.cloud.aiplatform_v1.types import EncryptionSpec
6566
from google.protobuf import timestamp_pb2
6667

6768
from google.cloud.aiplatform_v1.types import content
@@ -83,6 +84,9 @@
8384
TEST_WEAVIATE_API_KEY_SECRET_VERSION = (
8485
"projects/test-project/secrets/test-secret/versions/1"
8586
)
87+
TEST_ENCRYPTION_SPEC = EncryptionSpec(
88+
kms_key_name="projects/test-project/locations/us-central1/keyRings/test-key-ring/cryptoKeys/test-key"
89+
)
8690
TEST_PINECONE_INDEX_NAME = "test-pinecone-index"
8791
TEST_PINECONE_API_KEY_SECRET_VERSION = (
8892
"projects/test-project/secrets/test-secret/versions/1"
@@ -106,6 +110,14 @@
106110
TEST_GAPIC_RAG_CORPUS.vector_db_config.rag_embedding_model_config.vertex_prediction_endpoint.endpoint = "projects/{}/locations/{}/publishers/google/models/textembedding-gecko".format(
107111
TEST_PROJECT, TEST_REGION
108112
)
113+
TEST_GAPIC_CMEK_RAG_CORPUS = GapicRagCorpus(
114+
name=TEST_RAG_CORPUS_RESOURCE_NAME,
115+
display_name=TEST_CORPUS_DISPLAY_NAME,
116+
description=TEST_CORPUS_DISCRIPTION,
117+
encryption_spec=EncryptionSpec(
118+
kms_key_name="projects/test-project/locations/us-central1/keyRings/test-key-ring/cryptoKeys/test-key"
119+
),
120+
)
109121
TEST_GAPIC_RAG_CORPUS_VERTEX_VECTOR_SEARCH = GapicRagCorpus(
110122
name=TEST_RAG_CORPUS_RESOURCE_NAME,
111123
display_name=TEST_CORPUS_DISPLAY_NAME,
@@ -148,6 +160,14 @@
148160
display_name=TEST_CORPUS_DISPLAY_NAME,
149161
backend_config=TEST_BACKEND_CONFIG_EMBEDDING_MODEL_CONFIG,
150162
)
163+
TEST_CMEK_RAG_CORPUS = RagCorpus(
164+
name=TEST_RAG_CORPUS_RESOURCE_NAME,
165+
display_name=TEST_CORPUS_DISPLAY_NAME,
166+
description=TEST_CORPUS_DISCRIPTION,
167+
encryption_spec=EncryptionSpec(
168+
kms_key_name="projects/test-project/locations/us-central1/keyRings/test-key-ring/cryptoKeys/test-key"
169+
),
170+
)
151171
TEST_BACKEND_CONFIG_PINECONE_CONFIG = RagVectorDbConfig(
152172
vector_db=TEST_PINECONE_CONFIG,
153173
)

‎tests/unit/vertex_rag/test_rag_data.py

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,21 @@ def create_rag_corpus_mock():
5353
yield create_rag_corpus_mock
5454

5555

56+
@pytest.fixture
57+
def create_rag_corpus_mock_cmek():
58+
with mock.patch.object(
59+
VertexRagDataServiceClient,
60+
"create_rag_corpus",
61+
) as create_rag_corpus_mock_cmek:
62+
create_rag_corpus_lro_mock = mock.Mock(ga_operation.Operation)
63+
create_rag_corpus_lro_mock.done.return_value = True
64+
create_rag_corpus_lro_mock.result.return_value = (
65+
test_rag_constants.TEST_GAPIC_CMEK_RAG_CORPUS
66+
)
67+
create_rag_corpus_mock_cmek.return_value = create_rag_corpus_lro_mock
68+
yield create_rag_corpus_mock_cmek
69+
70+
5671
@pytest.fixture
5772
def create_rag_corpus_mock_vertex_vector_search():
5873
with mock.patch.object(
@@ -373,6 +388,15 @@ def test_create_corpus_vertex_vector_search_success(self):
373388
rag_corpus, test_rag_constants.TEST_RAG_CORPUS_VERTEX_VECTOR_SEARCH
374389
)
375390

391+
@pytest.mark.usefixtures("create_rag_corpus_mock_cmek")
392+
def test_create_corpus_cmek_success(self):
393+
rag_corpus = rag.create_corpus(
394+
display_name=test_rag_constants.TEST_CORPUS_DISPLAY_NAME,
395+
encryption_spec=test_rag_constants.TEST_ENCRYPTION_SPEC,
396+
)
397+
398+
rag_corpus_eq(rag_corpus, test_rag_constants.TEST_CMEK_RAG_CORPUS)
399+
376400
@pytest.mark.usefixtures("create_rag_corpus_mock_pinecone")
377401
def test_create_corpus_pinecone_success(self):
378402
rag_corpus = rag.create_corpus(

‎vertexai/rag/rag_data.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
from vertexai.rag.utils import (
4747
_gapic_utils,
4848
)
49+
from google.cloud.aiplatform_v1.types import EncryptionSpec
4950
from vertexai.rag.utils.resources import (
5051
JiraSource,
5152
LayoutParserConfig,
@@ -71,6 +72,7 @@ def create_corpus(
7172
None,
7273
]
7374
] = None,
75+
encryption_spec: Optional[EncryptionSpec] = None,
7476
) -> RagCorpus:
7577
"""Creates a new RagCorpus resource.
7678
@@ -96,6 +98,7 @@ def create_corpus(
9698
specified.
9799
backend_config: The backend config of the RagCorpus, specifying a
98100
data store and/or embedding model.
101+
encryption_spec: The encryption spec of the RagCorpus.
99102
Returns:
100103
RagCorpus.
101104
Raises:
@@ -124,6 +127,12 @@ def create_corpus(
124127
rag_corpus=rag_corpus,
125128
)
126129

130+
if encryption_spec:
131+
_gapic_utils.set_encryption_spec(
132+
encryption_spec=encryption_spec,
133+
rag_corpus=rag_corpus,
134+
)
135+
127136
request = CreateRagCorpusRequest(
128137
parent=parent,
129138
rag_corpus=rag_corpus,

‎vertexai/rag/utils/_gapic_utils.py

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import re
1818
from typing import Any, Dict, Optional, Sequence, Union
1919
from google.cloud.aiplatform_v1.types import api_auth
20+
from google.cloud.aiplatform_v1.types import EncryptionSpec
2021
from google.cloud.aiplatform_v1 import (
2122
RagEmbeddingModelConfig as GapicRagEmbeddingModelConfig,
2223
GoogleDriveSource,
@@ -203,6 +204,7 @@ def convert_gapic_to_rag_corpus(gapic_rag_corpus: GapicRagCorpus) -> RagCorpus:
203204
backend_config=convert_gapic_to_backend_config(
204205
gapic_rag_corpus.vector_db_config
205206
),
207+
encryption_spec=gapic_rag_corpus.encryption_spec,
206208
)
207209
return rag_corpus
208210

@@ -223,6 +225,7 @@ def convert_gapic_to_rag_corpus_no_embedding_model_config(
223225
backend_config=convert_gapic_to_backend_config(
224226
rag_vector_db_config_no_embedding_model_config
225227
),
228+
encryption_spec=gapic_rag_corpus.encryption_spec,
226229
)
227230
return rag_corpus
228231

@@ -660,6 +663,28 @@ def set_backend_config(
660663
)
661664

662665

666+
def set_encryption_spec(
667+
encryption_spec: EncryptionSpec,
668+
rag_corpus: GapicRagCorpus,
669+
) -> None:
670+
"""Sets the encryption spec for the rag corpus."""
671+
# Raises value error if encryption_spec.kms_key_name is None or empty,
672+
if encryption_spec.kms_key_name is None or not encryption_spec.kms_key_name:
673+
raise ValueError("kms_key_name must be set if encryption_spec is set.")
674+
675+
# Raises value error if encryption_spec.kms_key_name is not a valid KMS key name.
676+
if not re.match(
677+
r"^projects/(?P<project>.+?)/locations/(?P<location>.+?)/keyRings/(?P<key_ring>.+?)/cryptoKeys/(?P<crypto_key>.+?)$",
678+
encryption_spec.kms_key_name,
679+
):
680+
raise ValueError(
681+
"kms_key_name must be of the format "
682+
"`projects/{project}/locations/{location}/keyRings/{key_ring}/cryptoKeys/{crypto_key}`"
683+
)
684+
685+
rag_corpus.encryption_spec = encryption_spec
686+
687+
663688
def set_vertex_ai_search_config(
664689
vertex_ai_search_config: VertexAiSearchConfig,
665690
rag_corpus: GapicRagCorpus,

‎vertexai/rag/utils/resources.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
from typing import List, Optional, Sequence, Union
2020

2121
from google.protobuf import timestamp_pb2
22+
from google.cloud.aiplatform_v1.types import EncryptionSpec
2223

2324

2425
@dataclasses.dataclass
@@ -190,6 +191,7 @@ class RagCorpus:
190191
vertex_ai_search_config: The Vertex AI Search config of the RagCorpus.
191192
backend_config: The backend config of the RagCorpus. It can be a data
192193
store and/or retrieval engine.
194+
encryption_spec: The encryption spec of the RagCorpus. Immutable.
193195
"""
194196

195197
name: Optional[str] = None
@@ -202,6 +204,7 @@ class RagCorpus:
202204
None,
203205
]
204206
] = None
207+
encryption_spec: Optional[EncryptionSpec] = None
205208

206209

207210
@dataclasses.dataclass

0 commit comments

Comments
 (0)
0