fix: GenAI SDK (prompts) - Fix bug where passing encryption_spec to prompts.create raised an error

PiperOrigin-RevId: 821762191

Author: sararob

tests/unit/vertexai/genai/replays/test_create_prompt.py

@@ -17,9 +17,9 @@
 from tests.unit.vertexai.genai.replays import pytest_helper
 from vertexai._genai import types
 from google.genai import types as genai_types
-
 import pytest
 
+
 TEST_PROMPT_DATASET_ID = "8005484238453342208"
 TEST_VARIABLES = [
     {"name": genai_types.Part(text="Alice")},
@@ -289,6 +289,36 @@ def test_create_with_file_data(client):
     assert contents[0] == prompt_resource.prompt_data.contents[0]
 
 
+def test_create_with_encryption_spec(client):
+    encryption_spec = genai_types.EncryptionSpec(
+        kms_key_name="projects/vertex-sdk-dev/locations/us-central1/keyRings/my-key-ring/cryptoKeys/my-key",
+    )
+    config = types.CreatePromptConfig(
+        prompt_display_name="my_prompt_with_encryption_spec",
+        encryption_spec=encryption_spec,
+    )
+    prompt_resource = client.prompts.create(
+        prompt=TEST_PROMPT,
+        config=config,
+    )
+    assert isinstance(prompt_resource, types.Prompt)
+    assert isinstance(prompt_resource.dataset, types.Dataset)
+
+    # Create a version on a prompt with an encryption spec.
+    new_prompt = TEST_PROMPT.model_copy(deep=True)
+    new_prompt.prompt_data.contents[0].parts[0].text = "Is this Alice?"
+    prompt_version_resource = client.prompts.create_version(
+        prompt_id=prompt_resource.prompt_id,
+        prompt=new_prompt,
+        config=types.CreatePromptVersionConfig(
+            version_display_name="my_version_existing_dataset",
+        ),
+    )
+    assert isinstance(prompt_version_resource, types.Prompt)
+    assert isinstance(prompt_version_resource.dataset, types.Dataset)
+    assert isinstance(prompt_version_resource.dataset_version, types.DatasetVersion)
+
+
 pytestmark = pytest_helper.setup(
     file=__file__,
     globals_for_file=globals(),

vertexai/_genai/_prompt_management_utils.py

@@ -14,7 +14,7 @@
 #
 """Utility functions for prompt management."""
 
-from typing import Optional, Union
+from typing import Optional
 
 from google.genai import types as genai_types
 
@@ -123,20 +123,8 @@ def _create_prompt_from_dataset_metadata(
 
 def _raise_for_invalid_prompt(
     prompt: types.Prompt,
-    config: Optional[
-        Union[types.CreatePromptConfig, types.CreatePromptVersionConfig]
-    ] = None,
 ) -> None:
 
-    if (
-        isinstance(config, types.CreatePromptConfig)
-        and config.encryption_spec
-        and config.prompt_id
-    ):
-        raise ValueError(
-            "Encryption spec can only be used for creating new prompts, not for creating new prompt versions."
-        )
-
     if not prompt.prompt_data:
         raise ValueError("Prompt data must be provided.")
     if not prompt.prompt_data.contents:

vertexai/_genai/prompts.py

@@ -934,7 +934,7 @@ def create(
         elif not config:
             config = types.CreatePromptConfig()
 
-        _prompt_management_utils._raise_for_invalid_prompt(prompt, config)
+        _prompt_management_utils._raise_for_invalid_prompt(prompt)
 
         prompt_metadata = _prompt_management_utils._create_dataset_metadata_from_prompt(
             prompt,
@@ -1001,7 +1001,7 @@ def create_version(
         elif not config:
             config = types.CreatePromptVersionConfig()
 
-        _prompt_management_utils._raise_for_invalid_prompt(prompt, config)
+        _prompt_management_utils._raise_for_invalid_prompt(prompt)
 
         if config and config.version_display_name:
             version_name = config.version_display_name
@@ -2093,7 +2093,7 @@ async def create(
         elif not config:
             config = types.CreatePromptConfig()
 
-        _prompt_management_utils._raise_for_invalid_prompt(prompt, config)
+        _prompt_management_utils._raise_for_invalid_prompt(prompt)
 
         prompt_metadata = _prompt_management_utils._create_dataset_metadata_from_prompt(
             prompt,
@@ -2160,7 +2160,7 @@ async def create_version(
         elif not config:
             config = types.CreatePromptVersionConfig()
 
-        _prompt_management_utils._raise_for_invalid_prompt(prompt, config)
+        _prompt_management_utils._raise_for_invalid_prompt(prompt)
 
         if config and config.version_display_name:
             version_name = config.version_display_name