feat: support customer-managed encryption key in cached content

PiperOrigin-RevId: 759685739

Author: qiaodev

google/genai/caches.py

@@ -518,6 +518,9 @@ def _CreateCachedContentConfig_to_mldev(
         ),
     )
 
+  if getv(from_object, ['kms_key_name']) is not None:
+    raise ValueError('kms_key_name parameter is not supported in Gemini API.')
+
   return to_object
 
 
@@ -1171,6 +1174,13 @@ def _CreateCachedContentConfig_to_vertex(
         ),
     )
 
+  if getv(from_object, ['kms_key_name']) is not None:
+    setv(
+        parent_object,
+        ['encryption_spec', 'kmsKeyName'],
+        getv(from_object, ['kms_key_name']),
+    )
+
   return to_object
 
 

google/genai/tests/caches/test_create.py

@@ -93,6 +93,13 @@
     'models/gemini-1.5-pro-001'
 )
 
+_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI_CMEK = deepcopy(
+    _CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI
+)
+_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI_CMEK.config.kms_key_name = (
+    'projects/test-project/locations/us-central1/keyRings/test-keyring/cryptoKeys/test-key'
+)
+
 if sys.version_info >= (3, 11):
   _EXPIRE_TIME = datetime.datetime.fromisoformat('2024-12-20T00:00:00Z')
 else:
@@ -137,6 +144,12 @@
         exception_if_mldev='INVALID_ARGUMENT',
         parameters=_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI,
     ),
+    pytest_helper.TestTableItem(
+        name='test_caches_create_with_gcs_uri_cmek',
+        exception_if_mldev='not supported',
+        exception_if_vertex='INVALID_ARGUMENT',  # The key is invalid.
+        parameters=_CREATE_CACHED_CONTENT_PARAMETERS_GCS_URI_CMEK,
+    ),
     pytest_helper.TestTableItem(
         name='test_caches_create_with_gcs_uri_expire_time',
         exception_if_mldev='INVALID_ARGUMENT',

google/genai/types.py

@@ -8237,6 +8237,18 @@ class CreateCachedContentConfig(_common.BaseModel):
       description="""Configuration for the tools to use. This config is shared for all tools.
       """,
   )
+  kms_key_name: Optional[str] = Field(
+      default=None,
+      description="""The Cloud KMS resource identifier of the customer managed
+      encryption key used to protect a resource.
+      The key needs to be in the same region as where the compute resource is
+      created. 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}
+      """,
+  )
 
 
 class CreateCachedContentConfigDict(TypedDict, total=False):
@@ -8271,6 +8283,17 @@ class CreateCachedContentConfigDict(TypedDict, total=False):
   """Configuration for the tools to use. This config is shared for all tools.
       """
 
+  kms_key_name: Optional[str]
+  """The Cloud KMS resource identifier of the customer managed
+      encryption key used to protect a resource.
+      The key needs to be in the same region as where the compute resource is
+      created. 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}
+      """
+
 
 CreateCachedContentConfigOrDict = Union[
     CreateCachedContentConfig, CreateCachedContentConfigDict