feat: allow users to configure max wait time for prompt management queries

PiperOrigin-RevId: 944229591

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)
 
@@ -1565,6 +1565,9 @@ def create(
         dataset_resource_name = self._wait_for_operation(
             operation=create_prompt_dataset_operation,
             timeout=config.timeout if config else 90,
+            max_wait_time=(
+                config.max_wait_time if config and config.max_wait_time else 60
+            ),
         )
         dataset_id = dataset_resource_name.split("/")[-1]
 
@@ -1637,6 +1640,9 @@ def create_version(
         dataset_resource_name = self._wait_for_operation(
             operation=create_prompt_dataset_operation,
             timeout=config.timeout if config else 90,
+            max_wait_time=(
+                config.max_wait_time if config and config.max_wait_time else 60
+            ),
         )
         dataset_id = dataset_resource_name.split("/")[-1]
 
@@ -1661,6 +1667,9 @@ def create_version(
         dataset_version_resource_name = self._wait_for_operation(
             operation=create_dataset_version_operation,
             timeout=config.timeout if config else 90,
+            max_wait_time=(
+                config.max_wait_time if config and config.max_wait_time else 60
+            ),
         )
 
         # Step 4: Get the dataset version resource and return it with the prompt
@@ -1679,12 +1688,14 @@ def _wait_for_operation(
         self,
         operation: types.DatasetOperation,
         timeout: int,
+        max_wait_time: int = 60,
     ) -> str:
         """Waits for a dataset operation to complete.
 
         Args:
           operation: The dataset operation to wait for.
           timeout: The maximum time to wait for the operation to complete.
+          max_wait_time: The maximum interval between polling requests in seconds.
 
         Returns:
           The name of the Dataset resource from the operation result.
@@ -1706,7 +1717,6 @@ def _wait_for_operation(
         start_time = time.time()
         sleep_duration = 5
         wait_multiplier = 2
-        max_wait_time = 60
         previous_time = time.time()
 
         while not done:
@@ -1715,11 +1725,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 +1734,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
@@ -1923,6 +1936,7 @@ def _wait_for_project_operation(
         self,
         operation: genai_types.ProjectOperation,
         timeout: int,
+        max_wait_time: int = 60,
     ) -> None:
         """Waits for a dataset deletion operation to complete.
 
@@ -1931,6 +1945,7 @@ def _wait_for_project_operation(
         Args:
           operation: The project operation to wait for.
           timeout: The maximum time to wait for the operation to complete.
+          max_wait_time: The maximum interval between polling requests in seconds.
         Raises:
           TimeoutError: If the operation does not complete within the timeout.
           ValueError: If the operation fails.
@@ -1940,19 +1955,13 @@ def _wait_for_project_operation(
         start_time = time.time()
         sleep_duration = 5
         wait_multiplier = 2
-        max_wait_time = 60
         previous_time = time.time()
         while not done:
             if (time.time() - start_time) > timeout:
                 raise TimeoutError(
                     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 +1970,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}")
 
@@ -1985,7 +2002,11 @@ def delete(
             config=config,
         )
         self._wait_for_project_operation(
-            operation=delete_prompt_operation, timeout=config.timeout if config else 90
+            operation=delete_prompt_operation,
+            timeout=config.timeout if config and config.timeout else 90,
+            max_wait_time=(
+                config.max_wait_time if config and config.max_wait_time else 60
+            ),
         )
         logger.info(f"Deleted prompt with id: {prompt_id}")
 
@@ -2013,7 +2034,11 @@ def delete_version(
         )
 
         self._wait_for_project_operation(
-            operation=delete_version_operation, timeout=config.timeout if config else 90
+            operation=delete_version_operation,
+            timeout=config.timeout if config and config.timeout else 90,
+            max_wait_time=(
+                config.max_wait_time if config and config.max_wait_time else 60
+            ),
         )
         logger.info(
             f"Deleted prompt version {version_id} from prompt with id: {prompt_id}"
@@ -2043,7 +2068,10 @@ def restore_version(
         )
         self._wait_for_project_operation(
             operation=restore_prompt_operation,
-            timeout=90,
+            timeout=config.timeout if config and config.timeout else 90,
+            max_wait_time=(
+                config.max_wait_time if config and config.max_wait_time else 60
+            ),
         )
         dataset_version_resource = self._get_dataset_version_resource(
             dataset_id=prompt_id,
@@ -2369,7 +2397,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
@@ -2401,6 +2429,9 @@ def update(
         dataset_version_resource_name = self._wait_for_operation(
             operation=create_dataset_version_operation,
             timeout=config.timeout if config else 90,
+            max_wait_time=(
+                config.max_wait_time if config and config.max_wait_time else 60
+            ),
         )
         dataset_version_id = dataset_version_resource_name.split("/")[-1]
 
@@ -3159,7 +3190,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)
@@ -3559,6 +3590,9 @@ async def create(
         dataset_resource_name = await self._wait_for_operation(
             operation=create_prompt_dataset_operation,
             timeout=config.timeout if config else 90,
+            max_wait_time=(
+                config.max_wait_time if config and config.max_wait_time else 60
+            ),
         )
         dataset_id = dataset_resource_name.split("/")[-1]
 
@@ -3630,6 +3664,9 @@ async def create_version(
         dataset_resource_name = await self._wait_for_operation(
             operation=create_prompt_dataset_operation,
             timeout=config.timeout if config else 90,
+            max_wait_time=(
+                config.max_wait_time if config and config.max_wait_time else 60
+            ),
         )
         dataset_id = dataset_resource_name.split("/")[-1]
 
@@ -3654,6 +3691,9 @@ async def create_version(
         dataset_version_resource_name = await self._wait_for_operation(
             operation=create_dataset_version_operation,
             timeout=config.timeout if config else 90,
+            max_wait_time=(
+                config.max_wait_time if config and config.max_wait_time else 60
+            ),
         )
 
         # Step 4: Get the dataset version resource and return it with the prompt
@@ -3709,7 +3749,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
@@ -3741,6 +3781,9 @@ async def update(
         dataset_version_resource_name = await self._wait_for_operation(
             operation=create_dataset_version_operation,
             timeout=config.timeout if config else 90,
+            max_wait_time=(
+                config.max_wait_time if config and config.max_wait_time else 60
+            ),
         )
         dataset_version_id = dataset_version_resource_name.split("/")[-1]
 
@@ -3760,12 +3803,14 @@ async def _wait_for_operation(
         self,
         operation: types.DatasetOperation,
         timeout: int,
+        max_wait_time: int = 60,
     ) -> str:
         """Waits for a dataset operation to complete.
 
         Args:
           operation: The dataset operation to wait for.
           timeout: The maximum time to wait for the operation to complete.
+          max_wait_time: The maximum interval between polling requests in seconds.
 
         Returns:
           The name of the Dataset resource from the operation result.
@@ -3787,7 +3832,6 @@ async def _wait_for_operation(
         start_time = time.time()
         sleep_duration = 5
         wait_multiplier = 2
-        max_wait_time = 60
         previous_time = time.time()
 
         while not done:
@@ -3885,6 +3929,7 @@ async def _wait_for_project_operation(
         self,
         operation: genai_types.ProjectOperation,
         timeout: int,
+        max_wait_time: int = 60,
     ) -> None:
         """Waits for a dataset deletion operation to complete.
 
@@ -3893,6 +3938,7 @@ async def _wait_for_project_operation(
         Args:
           operation: The project operation to wait for.
           timeout: The maximum time to wait for the operation to complete.
+          max_wait_time: The maximum interval between polling requests in seconds.
         Raises:
           TimeoutError: If the operation does not complete within the timeout.
           ValueError: If the operation fails.
@@ -3902,19 +3948,13 @@ async def _wait_for_project_operation(
         start_time = time.time()
         sleep_duration = 5
         wait_multiplier = 2
-        max_wait_time = 60
         previous_time = time.time()
         while not done:
             if (time.time() - start_time) > timeout:
                 raise TimeoutError(
                     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 +3963,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}")
 
@@ -3947,7 +3995,11 @@ async def delete(
             config=config,
         )
         await self._wait_for_project_operation(
-            operation=delete_prompt_operation, timeout=config.timeout if config else 90
+            operation=delete_prompt_operation,
+            timeout=config.timeout if config and config.timeout else 90,
+            max_wait_time=(
+                config.max_wait_time if config and config.max_wait_time else 60
+            ),
         )
         logger.info(f"Deleted prompt with id: {prompt_id}")
 
@@ -3975,7 +4027,11 @@ async def delete_version(
         )
 
         await self._wait_for_project_operation(
-            operation=delete_version_operation, timeout=config.timeout if config else 90
+            operation=delete_version_operation,
+            timeout=config.timeout if config and config.timeout else 90,
+            max_wait_time=(
+                config.max_wait_time if config and config.max_wait_time else 60
+            ),
         )
         logger.info(
             f"Deleted prompt version {version_id} from prompt with id: {prompt_id}"
@@ -4114,7 +4170,10 @@ async def restore_version(
         )
         await self._wait_for_project_operation(
             operation=restore_prompt_operation,
-            timeout=90,
+            timeout=config.timeout if config and config.timeout else 90,
+            max_wait_time=(
+                config.max_wait_time if config and config.max_wait_time else 60
+            ),
         )
         dataset_version_resource = await self._get_dataset_version_resource(
             dataset_id=prompt_id,