feat: allow customers to set kms_key_name for context caching create method

vertex-sdk-bot · copybara-github · commit c551cc63cc13 · 2025-05-13T12:40:20.000-07:00
PiperOrigin-RevId: 758332266
diff --git a/tests/unit/vertexai/test_caching.py b/tests/unit/vertexai/test_caching.py
@@ -216,6 +216,34 @@ def test_create_with_real_payload(
         )
         assert cache.display_name == _TEST_DISPLAY_NAME
 
+    def test_create_with_encrypted_real_payload(
+        self, mock_create_cached_content, mock_get_cached_content
+    ):
+        cache = CachedContent.create(
+            model_name="model-name",
+            system_instruction=GapicContent(
+                role="system", parts=[GapicPart(text="system instruction")]
+            ),
+            tools=[],
+            tool_config=GapicToolConfig(),
+            contents=[GapicContent(role="user")],
+            ttl=datetime.timedelta(days=1),
+            display_name=_TEST_DISPLAY_NAME,
+            kms_key_name=f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/keyRings/test-key-ring/cryptoKeys/test-key",
+        )
+
+        # parent is automantically set to align with the current project and location.
+        # _CREATED_CONTENT_ID is from the mock
+        assert cache.resource_name == (
+            f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/cachedContents/{_CREATED_CONTENT_ID}"
+        )
+        assert cache.name == _CREATED_CONTENT_ID
+        assert (
+            cache.model_name
+            == f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/publishers/google/models/model-name"
+        )
+        assert cache.display_name == _TEST_DISPLAY_NAME
+
     def test_create_with_real_payload_and_wrapped_type(
         self, mock_create_cached_content, mock_get_cached_content
     ):
diff --git a/vertexai/caching/_caching.py b/vertexai/caching/_caching.py
@@ -39,6 +39,7 @@
     UpdateCachedContentRequest,
 )
 from google.cloud.aiplatform_v1 import types as types_v1
+from google.cloud.aiplatform_v1beta1.types import EncryptionSpec
 from vertexai.generative_models import _generative_models
 from vertexai.generative_models._generative_models import (
     Content,
@@ -60,6 +61,7 @@ def _prepare_create_request(
     expire_time: Optional[datetime.datetime] = None,
     ttl: Optional[datetime.timedelta] = None,
     display_name: Optional[str] = None,
+    kms_key_name: Optional[str] = None,
 ) -> CreateCachedContentRequest:
     """Prepares the request create_cached_content RPC."""
     (
@@ -103,6 +105,7 @@ def _prepare_create_request(
             expire_time=expire_time,
             ttl=ttl,
             display_name=display_name,
+            encryption_spec=EncryptionSpec(kms_key_name=kms_key_name),
         ),
     )
     serialized_message_v1beta1 = type(request_v1beta1).serialize(request_v1beta1)
@@ -175,6 +178,7 @@ def create(
         expire_time: Optional[datetime.datetime] = None,
         ttl: Optional[datetime.timedelta] = None,
         display_name: Optional[str] = None,
+        kms_key_name: Optional[str] = None,
     ) -> "CachedContent":
         """Creates a new cached content through the gen ai cache service.
 
@@ -213,6 +217,14 @@ def create(
                 default TTL on the API side will be used (currently 1 hour).
             display_name:
                 The user-generated meaningful display name of the cached content.
+            kms_key_name:
+                Optional. Customer-managed encryption key. See
+                https://cloud.google.com/vertex-ai/docs/general/cmek for more
+                details. If this is set, then all created CachedContent objects
+                will be encrypted with the provided encryption key.
+                Allowed formats:
+
+                projects/{project}/locations/{location}/keyRings/{key_ring}/cryptoKeys/{crypto_key}
         Returns:
             A CachedContent object with only name and model_name specified.
         Raises:
@@ -237,6 +249,7 @@ def create(
             expire_time=expire_time,
             ttl=ttl,
             display_name=display_name,
+            kms_key_name=kms_key_name,
         )
         client = cls._instantiate_client(location=location)
         cached_content_resource = client.create_cached_content(request)
