feat: Support Prompt Management in the JS GenAI Modules

PiperOrigin-RevId: 903243258

Author: speedstorm1

agentplatform/_genai/prompts.py

@@ -1174,7 +1174,7 @@ def _restore_version(
         request_dict = _common.convert_to_dict(request_dict)
         request_dict = _common.encode_unserializable_types(request_dict)
 
-        response = self._api_client.request("get", path, request_dict, http_options)
+        response = self._api_client.request("post", path, request_dict, http_options)
 
         response_dict = {} if not response.body else json.loads(response.body)
 
@@ -1715,11 +1715,6 @@ def _wait_for_operation(
                     "Create prompt operation did not complete within the"
                     f" specified timeout of {timeout} seconds."
                 )
-            current_time = time.time()
-            if current_time - previous_time >= sleep_duration:
-                sleep_duration = min(sleep_duration * wait_multiplier, max_wait_time)
-                previous_time = current_time
-            time.sleep(sleep_duration)
             prompt_dataset_operation = self._get_dataset_operation(
                 dataset_id=dataset_id,
                 operation_id=operation_id,
@@ -1729,6 +1724,14 @@ def _wait_for_operation(
                 if hasattr(prompt_dataset_operation, "done")
                 else False
             )
+            if not done:
+                current_time = time.time()
+                if current_time - previous_time >= sleep_duration:
+                    sleep_duration = min(
+                        sleep_duration * wait_multiplier, max_wait_time
+                    )
+                    previous_time = current_time
+                time.sleep(sleep_duration)
         if (
             not prompt_dataset_operation
             or prompt_dataset_operation.response is None
@@ -1948,11 +1951,6 @@ def _wait_for_project_operation(
                     f"Delete operation did not complete within the"
                     f" specified timeout of {timeout} seconds."
                 )
-            current_time = time.time()
-            if current_time - previous_time >= sleep_duration:
-                sleep_duration = min(sleep_duration * wait_multiplier, max_wait_time)
-                previous_time = current_time
-            time.sleep(sleep_duration)
             operations_module = operations.Operations(api_client_=self._api_client)
 
             if operation.name is None:
@@ -1961,6 +1959,14 @@ def _wait_for_project_operation(
                 operation_id=operation.name.split("/")[-1],
             )
             done = (operation.done or False) if hasattr(operation, "done") else False
+            if not done:
+                current_time = time.time()
+                if current_time - previous_time >= sleep_duration:
+                    sleep_duration = min(
+                        sleep_duration * wait_multiplier, max_wait_time
+                    )
+                    previous_time = current_time
+                time.sleep(sleep_duration)
         if hasattr(operation, "error") and operation.error is not None:
             raise ValueError(f"Error in delete operation: {operation.error}")
 
@@ -2369,7 +2375,7 @@ def update(
 
         # Step 1: Update the dataset resource for the prompt and wait for the operation to complete.
         updated_dataset_resource = self._update_dataset_resource(
-            name=f"projects/{self._api_client.project}/locations/{self._api_client.location}",
+            name=f"projects/{self._api_client.project}/locations/{self._api_client.location}/datasets/{prompt_id}",
             dataset_id=prompt_id,
             display_name=(
                 config.prompt_display_name
@@ -3159,7 +3165,7 @@ async def _restore_version(
         request_dict = _common.encode_unserializable_types(request_dict)
 
         response = await self._api_client.async_request(
-            "get", path, request_dict, http_options
+            "post", path, request_dict, http_options
         )
 
         response_dict = {} if not response.body else json.loads(response.body)
@@ -3709,7 +3715,7 @@ async def update(
 
         # Step 1: Update the dataset resource for the prompt and wait for the operation to complete.
         updated_dataset_resource = await self._update_dataset_resource(
-            name=f"projects/{self._api_client.project}/locations/{self._api_client.location}",
+            name=f"projects/{self._api_client.project}/locations/{self._api_client.location}/datasets/{prompt_id}",
             dataset_id=prompt_id,
             display_name=(
                 config.prompt_display_name
@@ -3910,11 +3916,6 @@ async def _wait_for_project_operation(
                     f"Delete operation did not complete within the"
                     f" specified timeout of {timeout} seconds."
                 )
-            current_time = time.time()
-            if current_time - previous_time >= sleep_duration:
-                sleep_duration = min(sleep_duration * wait_multiplier, max_wait_time)
-                previous_time = current_time
-            await asyncio.sleep(sleep_duration)
             operations_module = operations.AsyncOperations(api_client_=self._api_client)
 
             if operation.name is None:
@@ -3923,6 +3924,14 @@ async def _wait_for_project_operation(
                 operation_id=operation.name.split("/")[-1],
             )
             done = (operation.done or False) if hasattr(operation, "done") else False
+            if not done:
+                current_time = time.time()
+                if current_time - previous_time >= sleep_duration:
+                    sleep_duration = min(
+                        sleep_duration * wait_multiplier, max_wait_time
+                    )
+                    previous_time = current_time
+                await asyncio.sleep(sleep_duration)
         if hasattr(operation, "error") and operation.error is not None:
             raise ValueError(f"Error in delete operation: {operation.error}")