chore: GenAI SDK client - refactor prompt management operations to use Operation type

PiperOrigin-RevId: 804577329

Author: sararob

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

@@ -80,7 +80,6 @@
 
 def test_create_dataset(client):
     create_dataset_operation = client.prompt_management._create_dataset_resource(
-        config=types.CreateDatasetConfig(should_return_http_response=True),
         name="projects/vertex-sdk-dev/locations/us-central1",
         display_name="test display name",
         metadata_schema_uri="gs://google-cloud-aiplatform/schema/dataset/metadata/text_prompt_1.0.0.yaml",
@@ -123,8 +122,8 @@ def test_create_dataset(client):
         },
         model_reference="gemini-2.0-flash-001",
     )
-    assert isinstance(create_dataset_operation, types.CreateDatasetOperationMetadata)
-    assert create_dataset_operation.sdk_http_response.body is not None
+    assert isinstance(create_dataset_operation, types.DatasetOperation)
+    assert create_dataset_operation
 
 
 def test_create_dataset_version(client):
@@ -134,9 +133,7 @@ def test_create_dataset_version(client):
             display_name="my new version yay",
         )
     )
-    assert isinstance(
-        dataset_version_resource, types.CreateDatasetVersionOperationMetadata
-    )
+    assert isinstance(dataset_version_resource, types.DatasetOperation)
 
 
 def test_create_version_e2e(client):

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

@@ -15,16 +15,14 @@
 # pylint: disable=protected-access,bad-continuation,missing-function-docstring
 
 from tests.unit.vertexai.genai.replays import pytest_helper
-from vertexai._genai import types
 
 
 def test_get_dataset_operation(client):
     dataset_operation = client.prompt_management._get_dataset_operation(
-        config=types.GetDatasetOperationConfig(should_return_http_response=True),
         dataset_id="6550997480673116160",
         operation_id="5108504762664353792",
     )
-    assert dataset_operation.sdk_http_response.body is not None
+    assert dataset_operation.name is not None
 
 
 pytestmark = pytest_helper.setup(

vertexai/_genai/_prompt_management_utils.py

@@ -14,7 +14,6 @@
 #
 """Utility functions for prompt management."""
 
-import json
 from typing import Optional
 
 from google.genai import types as genai_types
@@ -78,19 +77,6 @@ def _create_dataset_metadata_from_prompt(
     return prompt_metadata
 
 
-def _check_operation_status_from_response_body(
-    operation_response: genai_types.HttpResponse,
-):
-    """Checks the operation status from the response body."""
-    if not hasattr(operation_response, "body"):
-        raise ValueError("Error creating prompt dataset resource.")
-    done = False
-    body = json.loads(operation_response.body)
-    if body.get("done"):
-        done = True
-    return done
-
-
 def _create_prompt_from_dataset_metadata(
     dataset: types.Dataset,
 ) -> types.Prompt: