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

PiperOrigin-RevId: 944229591

Author: speedstorm1

agentplatform/_genai/prompts.py

@@ -1564,7 +1564,12 @@ def create(
         )
         dataset_resource_name = self._wait_for_operation(
             operation=create_prompt_dataset_operation,
-            timeout=config.timeout if config else 90,
+            timeout=config.timeout if config and config.timeout is not None else 90,
+            max_wait_time=(
+                config.max_wait_time
+                if config and config.max_wait_time is not None
+                else 60
+            ),
         )
         dataset_id = dataset_resource_name.split("/")[-1]
 
@@ -1636,7 +1641,12 @@ def create_version(
         )
         dataset_resource_name = self._wait_for_operation(
             operation=create_prompt_dataset_operation,
-            timeout=config.timeout if config else 90,
+            timeout=config.timeout if config and config.timeout is not None else 90,
+            max_wait_time=(
+                config.max_wait_time
+                if config and config.max_wait_time is not None
+                else 60
+            ),
         )
         dataset_id = dataset_resource_name.split("/")[-1]
 
@@ -1660,7 +1670,12 @@ def create_version(
         )
         dataset_version_resource_name = self._wait_for_operation(
             operation=create_dataset_version_operation,
-            timeout=config.timeout if config else 90,
+            timeout=config.timeout if config and config.timeout is not None else 90,
+            max_wait_time=(
+                config.max_wait_time
+                if config and config.max_wait_time is not None
+                else 60
+            ),
         )
 
         # Step 4: Get the dataset version resource and return it with the prompt
@@ -1679,12 +1694,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 +1723,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:
@@ -1923,6 +1939,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 +1948,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,7 +1958,6 @@ 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:
@@ -1985,7 +2002,13 @@ 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 is not None else 90,
+            max_wait_time=(
+                config.max_wait_time
+                if config and config.max_wait_time is not None
+                else 60
+            ),
         )
         logger.info(f"Deleted prompt with id: {prompt_id}")
 
@@ -2013,7 +2036,13 @@ 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 is not None else 90,
+            max_wait_time=(
+                config.max_wait_time
+                if config and config.max_wait_time is not None
+                else 60
+            ),
         )
         logger.info(
             f"Deleted prompt version {version_id} from prompt with id: {prompt_id}"
@@ -2040,10 +2069,16 @@ def restore_version(
         restore_prompt_operation = self._restore_version(
             dataset_id=prompt_id,
             version_id=version_id,
+            config=config,
         )
         self._wait_for_project_operation(
             operation=restore_prompt_operation,
-            timeout=90,
+            timeout=config.timeout if config and config.timeout is not None else 90,
+            max_wait_time=(
+                config.max_wait_time
+                if config and config.max_wait_time is not None
+                else 60
+            ),
         )
         dataset_version_resource = self._get_dataset_version_resource(
             dataset_id=prompt_id,
@@ -2400,7 +2435,12 @@ def update(
         )
         dataset_version_resource_name = self._wait_for_operation(
             operation=create_dataset_version_operation,
-            timeout=config.timeout if config else 90,
+            timeout=config.timeout if config and config.timeout is not None else 90,
+            max_wait_time=(
+                config.max_wait_time
+                if config and config.max_wait_time is not None
+                else 60
+            ),
         )
         dataset_version_id = dataset_version_resource_name.split("/")[-1]
 
@@ -3558,7 +3598,12 @@ async def create(
         )
         dataset_resource_name = await self._wait_for_operation(
             operation=create_prompt_dataset_operation,
-            timeout=config.timeout if config else 90,
+            timeout=config.timeout if config and config.timeout is not None else 90,
+            max_wait_time=(
+                config.max_wait_time
+                if config and config.max_wait_time is not None
+                else 60
+            ),
         )
         dataset_id = dataset_resource_name.split("/")[-1]
 
@@ -3629,7 +3674,12 @@ async def create_version(
         )
         dataset_resource_name = await self._wait_for_operation(
             operation=create_prompt_dataset_operation,
-            timeout=config.timeout if config else 90,
+            timeout=config.timeout if config and config.timeout is not None else 90,
+            max_wait_time=(
+                config.max_wait_time
+                if config and config.max_wait_time is not None
+                else 60
+            ),
         )
         dataset_id = dataset_resource_name.split("/")[-1]
 
@@ -3653,7 +3703,12 @@ 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,
+            timeout=config.timeout if config and config.timeout is not None else 90,
+            max_wait_time=(
+                config.max_wait_time
+                if config and config.max_wait_time is not None
+                else 60
+            ),
         )
 
         # Step 4: Get the dataset version resource and return it with the prompt
@@ -3740,7 +3795,12 @@ async def update(
         )
         dataset_version_resource_name = await self._wait_for_operation(
             operation=create_dataset_version_operation,
-            timeout=config.timeout if config else 90,
+            timeout=config.timeout if config and config.timeout is not None else 90,
+            max_wait_time=(
+                config.max_wait_time
+                if config and config.max_wait_time is not None
+                else 60
+            ),
         )
         dataset_version_id = dataset_version_resource_name.split("/")[-1]
 
@@ -3760,12 +3820,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 +3849,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 +3946,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 +3955,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,7 +3965,6 @@ 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:
@@ -3947,7 +4009,13 @@ 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 is not None else 90,
+            max_wait_time=(
+                config.max_wait_time
+                if config and config.max_wait_time is not None
+                else 60
+            ),
         )
         logger.info(f"Deleted prompt with id: {prompt_id}")
 
@@ -3975,7 +4043,13 @@ 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 is not None else 90,
+            max_wait_time=(
+                config.max_wait_time
+                if config and config.max_wait_time is not None
+                else 60
+            ),
         )
         logger.info(
             f"Deleted prompt version {version_id} from prompt with id: {prompt_id}"
@@ -4111,10 +4185,16 @@ async def restore_version(
         restore_prompt_operation = await self._restore_version(
             dataset_id=prompt_id,
             version_id=version_id,
+            config=config,
         )
         await self._wait_for_project_operation(
             operation=restore_prompt_operation,
-            timeout=90,
+            timeout=config.timeout if config and config.timeout is not None else 90,
+            max_wait_time=(
+                config.max_wait_time
+                if config and config.max_wait_time is not None
+                else 60
+            ),
         )
         dataset_version_resource = await self._get_dataset_version_resource(
             dataset_id=prompt_id,