feat: GenAI SDK client - Update client.prompts.create to create a prompt and prompt version for a prompt. Deprecate client.prompts.create_version in favor of create()

PiperOrigin-RevId: 889955089

Author: sararob

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

@@ -147,6 +147,16 @@ def test_create(client):
     assert isinstance(prompt_resource.dataset, types.Dataset)
 
 
+def test_create_updated_e2e(client):
+    prompt_resource = client.prompts.create(
+        prompt=TEST_PROMPT,
+        config=TEST_CREATE_PROMPT_CONFIG,
+    )
+    assert isinstance(prompt_resource, types.Prompt)
+    assert isinstance(prompt_resource.dataset, types.Dataset)
+    assert isinstance(prompt_resource.dataset_version, types.DatasetVersion)
+
+
 def test_create_e2e(client):
     prompt_resource = client.prompts.create(
         prompt=TEST_PROMPT,
@@ -361,3 +371,14 @@ async def test_create_version_async(client):
         prompt_version_resource.prompt_data.contents[0].parts[0].text
         == "Is this Alice?"
     )
+
+
+@pytest.mark.asyncio
+async def test_create_updated_e2e_async(client):
+    prompt_resource = await client.aio.prompts.create(
+        prompt=TEST_PROMPT,
+        config=TEST_CREATE_PROMPT_CONFIG,
+    )
+    assert isinstance(prompt_resource, types.Prompt)
+    assert isinstance(prompt_resource.dataset, types.Dataset)
+    assert isinstance(prompt_resource.dataset_version, types.DatasetVersion)

vertexai/_genai/prompts.py

@@ -1305,6 +1305,31 @@ def create(
             name=dataset_id,
         )
         prompt._dataset = dataset_resource
+
+        # Step 3: Create the dataset version
+        create_dataset_version_operation = self._create_dataset_version_resource(
+            dataset_name=prompt._dataset.name.split("/")[-1],
+            display_name=(
+                config.version_display_name
+                if config and config.version_display_name is not None
+                else f"prompt_version_{time.strftime('%Y%m%d-%H%M%S')}"
+            ),
+        )
+        dataset_version_resource_name = self._wait_for_operation(
+            operation=create_dataset_version_operation,
+            timeout=config.timeout if config else 90,
+        )
+
+        # Step 4: Get the dataset version resource and return it with the prompt
+        dataset_version_resource = self._get_dataset_version_resource(
+            dataset_id=dataset_id,
+            dataset_version_id=dataset_version_resource_name.split("/")[-1],
+        )
+        prompt = _prompt_management_utils._create_prompt_from_dataset_metadata(
+            dataset_version_resource
+        )
+        prompt._dataset = dataset_resource
+        prompt._dataset_version = dataset_version_resource
         return prompt
 
     def create_version(
@@ -1328,6 +1353,9 @@ def create_version(
           A types.Prompt object representing the prompt with its associated
           Dataset and Dataset Version resources.
         """
+        logger.warning(
+            "The `prompts.create_version` method is deprecated. Please use the prompts.create` method instead. `prompts.create` creates a prompt dataset and dataset version resource for your prompt."
+        )
         if isinstance(prompt, dict):
             prompt = types.Prompt(**prompt)
         if isinstance(config, dict):
@@ -2925,6 +2953,31 @@ async def create(
             name=dataset_id,
         )
         prompt._dataset = dataset_resource
+
+        # Step 3: Create the dataset version
+        create_dataset_version_operation = await self._create_dataset_version_resource(
+            dataset_name=dataset_id,
+            display_name=(
+                config.version_display_name
+                if config and config.version_display_name is not None
+                else f"prompt_version_{time.strftime('%Y%m%d-%H%M%S')}"
+            ),
+        )
+        dataset_version_resource_name = await self._wait_for_operation(
+            operation=create_dataset_version_operation,
+            timeout=config.timeout if config else 90,
+        )
+
+        # Step 4: Get the dataset version resource and return it with the prompt
+        dataset_version_resource = await self._get_dataset_version_resource(
+            dataset_id=dataset_id,
+            dataset_version_id=dataset_version_resource_name.split("/")[-1],
+        )
+        prompt = _prompt_management_utils._create_prompt_from_dataset_metadata(
+            dataset_version_resource
+        )
+        prompt._dataset = dataset_resource
+        prompt._dataset_version = dataset_version_resource
         return prompt
 
     async def create_version(
@@ -2948,6 +3001,9 @@ async def create_version(
           A types.Prompt object representing the prompt with its associated
           Dataset and Dataset Version resources.
         """
+        logger.warning(
+            "The `prompts.create_version` method is deprecated. Please use the prompts.create` method instead. `prompts.create` creates a prompt dataset and dataset version resource for your prompt."
+        )
         if isinstance(prompt, dict):
             prompt = types.Prompt(**prompt)
         if isinstance(config, dict):

vertexai/_genai/types/common.py

@@ -16043,6 +16043,10 @@ class CreatePromptConfig(_common.BaseModel):
         default=None,
         description="""Customer-managed encryption key spec for a prompt dataset. If set, this prompt dataset and all sub-resources of this prompt dataset will be secured by this key.""",
     )
+    version_display_name: Optional[str] = Field(
+        default=None,
+        description="""The display name for the prompt version. If not set, a default name with a timestamp will be used.""",
+    )
 
 
 class CreatePromptConfigDict(TypedDict, total=False):
@@ -16060,6 +16064,9 @@ class CreatePromptConfigDict(TypedDict, total=False):
     encryption_spec: Optional[genai_types.EncryptionSpecDict]
     """Customer-managed encryption key spec for a prompt dataset. If set, this prompt dataset and all sub-resources of this prompt dataset will be secured by this key."""
 
+    version_display_name: Optional[str]
+    """The display name for the prompt version. If not set, a default name with a timestamp will be used."""
+
 
 CreatePromptConfigOrDict = Union[CreatePromptConfig, CreatePromptConfigDict]