Content-Length: 729697 | pFad | https://github.com/googleapis/python-aiplatform/commit/ad821c58f834981adb9759013248e7156f3706b9

D3 feat: Update preview `create_corpus` to accept `encryption_spec` in `… · googleapis/python-aiplatform@ad821c5 · GitHub
Skip to content

Commit ad821c5

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
feat: Update preview create_corpus to accept encryption_spec in rag_data.py
PiperOrigin-RevId: 761091381
1 parent e13f19f commit ad821c5

File tree

5 files changed

+85
-2
lines changed

5 files changed

+85
-2
lines changed

‎tests/unit/vertex_rag/test_rag_constants_preview.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
VertexAiSearchConfig as GapicVertexAiSearchConfig,
3939
)
4040
from google.cloud.aiplatform_v1beta1.types import api_auth
41+
from google.cloud.aiplatform_v1beta1.types import EncryptionSpec
4142
from vertexai.preview.rag import (
4243
ANN,
4344
Basic,
@@ -91,6 +92,9 @@
9192
TEST_WEAVIATE_API_KEY_SECRET_VERSION = (
9293
"projects/test-project/secrets/test-secret/versions/1"
9394
)
95+
TEST_ENCRYPTION_SPEC = EncryptionSpec(
96+
kms_key_name="projects/test-project/locations/us-central1/keyRings/test-key-ring/cryptoKeys/test-key"
97+
)
9498
TEST_WEAVIATE_CONFIG = Weaviate(
9599
weaviate_http_endpoint=TEST_WEAVIATE_HTTP_ENDPOINT,
96100
collection_name=TEST_WEAVIATE_COLLECTION_NAME,
@@ -133,6 +137,14 @@
133137
TEST_PROJECT, TEST_REGION
134138
)
135139
)
140+
TEST_GAPIC_CMEK_RAG_CORPUS = GapicRagCorpus(
141+
name=TEST_RAG_CORPUS_RESOURCE_NAME,
142+
display_name=TEST_CORPUS_DISPLAY_NAME,
143+
description=TEST_CORPUS_DISCRIPTION,
144+
encryption_spec=EncryptionSpec(
145+
kms_key_name="projects/test-project/locations/us-central1/keyRings/test-key-ring/cryptoKeys/test-key"
146+
),
147+
)
136148
TEST_GAPIC_RAG_CORPUS_WEAVIATE = GapicRagCorpus(
137149
name=TEST_RAG_CORPUS_RESOURCE_NAME,
138150
display_name=TEST_CORPUS_DISPLAY_NAME,
@@ -236,6 +248,14 @@
236248
description=TEST_CORPUS_DISCRIPTION,
237249
embedding_model_config=TEST_EMBEDDING_MODEL_CONFIG,
238250
)
251+
TEST_CMEK_RAG_CORPUS = RagCorpus(
252+
name=TEST_RAG_CORPUS_RESOURCE_NAME,
253+
display_name=TEST_CORPUS_DISPLAY_NAME,
254+
description=TEST_CORPUS_DISCRIPTION,
255+
encryption_spec=EncryptionSpec(
256+
kms_key_name="projects/test-project/locations/us-central1/keyRings/test-key-ring/cryptoKeys/test-key"
257+
),
258+
)
239259
TEST_RAG_CORPUS_WEAVIATE = RagCorpus(
240260
name=TEST_RAG_CORPUS_RESOURCE_NAME,
241261
display_name=TEST_CORPUS_DISPLAY_NAME,

‎tests/unit/vertex_rag/test_rag_data_preview.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_preview.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_backend():
5873
with mock.patch.object(
@@ -653,6 +668,15 @@ def test_create_corpus_backend_success(self):
653668

654669
rag_corpus_eq(rag_corpus, test_rag_constants_preview.TEST_RAG_CORPUS_BACKEND)
655670

671+
@pytest.mark.usefixtures("create_rag_corpus_mock_cmek")
672+
def test_create_corpus_cmek_success(self):
673+
rag_corpus = rag.create_corpus(
674+
display_name=test_rag_constants_preview.TEST_CORPUS_DISPLAY_NAME,
675+
encryption_spec=test_rag_constants_preview.TEST_ENCRYPTION_SPEC,
676+
)
677+
678+
rag_corpus_eq(rag_corpus, test_rag_constants_preview.TEST_CMEK_RAG_CORPUS)
679+
656680
@pytest.mark.usefixtures("create_rag_corpus_mock_weaviate")
657681
def test_create_corpus_weaviate_success(self):
658682
rag_corpus = rag.create_corpus(

‎vertexai/preview/rag/rag_data.py

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@
4141
ListRagCorporaPager,
4242
ListRagFilesPager,
4343
)
44+
from google.cloud.aiplatform_v1beta1.types import EncryptionSpec
4445
from vertexai.preview.rag.utils import (
4546
_gapic_utils,
4647
)
@@ -74,6 +75,7 @@ def create_corpus(
7475
] = None,
7576
vertex_ai_search_config: Optional[VertexAiSearchConfig] = None,
7677
backend_config: Optional[RagVectorDbConfig] = None,
78+
encryption_spec: Optional[EncryptionSpec] = None,
7779
) -> RagCorpus:
7880
"""Creates a new RagCorpus resource.
7981
@@ -91,8 +93,8 @@ def create_corpus(
9193
9294
Args:
9395
display_name: If not provided, SDK will create one. The display name of
94-
the RagCorpus. The name can be up to 128 characters long and can
95-
consist of any UTF-8 characters.
96+
the RagCorpus. The name can be up to 128 characters long and can consist
97+
of any UTF-8 characters.
9698
description: The description of the RagCorpus.
9799
embedding_model_config: The embedding model config.
98100
Note: Deprecated. Use backend_config instead.
@@ -104,6 +106,8 @@ def create_corpus(
104106
vertex_ai_search_config is specified.
105107
backend_config: The backend config of the RagCorpus. It can specify a
106108
Vector DB and/or the embedding model config.
109+
encryption_spec: The encryption spec of the RagCorpus.
110+
107111
Returns:
108112
RagCorpus.
109113
Raises:
@@ -156,6 +160,12 @@ def create_corpus(
156160
rag_corpus=rag_corpus,
157161
)
158162

163+
if encryption_spec:
164+
_gapic_utils.set_encryption_spec(
165+
encryption_spec=encryption_spec,
166+
rag_corpus=rag_corpus,
167+
)
168+
159169
request = CreateRagCorpusRequest(
160170
parent=parent,
161171
rag_corpus=rag_corpus,

‎vertexai/preview/rag/utils/_gapic_utils.py

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@
4141
VertexAiSearchConfig as GapicVertexAiSearchConfig,
4242
)
4343
from google.cloud.aiplatform_v1beta1.types import api_auth
44+
from google.cloud.aiplatform_v1beta1.types import EncryptionSpec
4445
from vertexai.preview.rag.utils.resources import (
4546
ANN,
4647
EmbeddingModelConfig,
@@ -327,6 +328,7 @@ def convert_gapic_to_rag_corpus(gapic_rag_corpus: GapicRagCorpus) -> RagCorpus:
327328
backend_config=convert_gapic_to_backend_config(
328329
gapic_rag_corpus.vector_db_config
329330
),
331+
encryption_spec=gapic_rag_corpus.encryption_spec,
330332
)
331333
return rag_corpus
332334

@@ -348,6 +350,7 @@ def convert_gapic_to_rag_corpus_no_embedding_model_config(
348350
backend_config=convert_gapic_to_backend_config(
349351
rag_vector_db_config_no_embedding_model_config
350352
),
353+
encryption_spec=gapic_rag_corpus.encryption_spec,
351354
)
352355
return rag_corpus
353356

@@ -728,6 +731,28 @@ def set_embedding_model_config(
728731
)
729732

730733

734+
def set_encryption_spec(
735+
encryption_spec: EncryptionSpec,
736+
rag_corpus: GapicRagCorpus,
737+
) -> None:
738+
"""Sets the encryption spec for the rag corpus."""
739+
# Raises value error if encryption_spec.kms_key_name is None or empty,
740+
if encryption_spec.kms_key_name is None or not encryption_spec.kms_key_name:
741+
raise ValueError("kms_key_name must be set if encryption_spec is set.")
742+
743+
# Raises value error if encryption_spec.kms_key_name is not a valid KMS key name.
744+
if not re.match(
745+
r"^projects/(?P<project>.+?)/locations/(?P<location>.+?)/keyRings/(?P<key_ring>.+?)/cryptoKeys/(?P<crypto_key>.+?)$",
746+
encryption_spec.kms_key_name,
747+
):
748+
raise ValueError(
749+
"kms_key_name must be of the format "
750+
"`projects/{project}/locations/{location}/keyRings/{key_ring}/cryptoKeys/{crypto_key}`"
751+
)
752+
753+
rag_corpus.encryption_spec = encryption_spec
754+
755+
731756
def set_vector_db(
732757
vector_db: Union[
733758
Weaviate, VertexFeatureStore, VertexVectorSearch, Pinecone, RagManagedDb, None

‎vertexai/preview/rag/utils/resources.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@
1818
import dataclasses
1919
from typing import List, Optional, Sequence, Union
2020

21+
from google.cloud.aiplatform_v1beta1.types import EncryptionSpec
22+
2123
from google.protobuf import timestamp_pb2
2224

2325
DEPRECATION_DATE = "June 2025"
@@ -270,6 +272,7 @@ class RagCorpus:
270272
vertex_ai_search_config: The Vertex AI Search config of the RagCorpus.
271273
backend_config: The backend config of the RagCorpus. It can specify a
272274
Vector DB and/or the embedding model config.
275+
encryption_spec: The encryption spec of the RagCorpus. Immutable.
273276
"""
274277

275278
name: Optional[str] = None
@@ -281,6 +284,7 @@ class RagCorpus:
281284
] = None
282285
vertex_ai_search_config: Optional[VertexAiSearchConfig] = None
283286
backend_config: Optional[RagVectorDbConfig] = None
287+
encryption_spec: Optional[EncryptionSpec] = None
284288

285289

286290
@dataclasses.dataclass

0 commit comments

Comments
 (0)








ApplySandwichStrip

pFad - (p)hone/(F)rame/(a)nonymizer/(d)eclutterfier!      Saves Data!


--- a PPN by Garber Painting Akron. With Image Size Reduction included!

Fetched URL: https://github.com/googleapis/python-aiplatform/commit/ad821c58f834981adb9759013248e7156f3706b9

Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy