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

FUTURE_COPYBARA_INTEGRATE_REVIEW=#7003 from googleapis:release-please--branches--main 5905f2d
PiperOrigin-RevId: 944229591

Author: speedstorm1

agentplatform/_genai/prompts.py

@@ -1564,7 +1564,10 @@ 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 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]
 
@@ -1636,7 +1639,10 @@ 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 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]
 
@@ -1660,7 +1666,10 @@ 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 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:
@@ -1923,6 +1933,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 +1942,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 +1952,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 +1996,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 +2028,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}"
@@ -2040,10 +2059,14 @@ 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 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,
@@ -2400,7 +2423,10 @@ 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 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]
 
@@ -3558,7 +3584,10 @@ 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 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]
 
@@ -3629,7 +3658,10 @@ 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 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]
 
@@ -3653,7 +3685,10 @@ 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 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
@@ -3740,7 +3775,10 @@ 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 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 +3798,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 +3827,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 +3924,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 +3933,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 +3943,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 +3987,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 +4019,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}"
@@ -4111,10 +4159,14 @@ 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 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,

agentplatform/_genai/types/common.py

@@ -21104,6 +21104,10 @@ class DeletePromptConfig(_common.BaseModel):
         default=90,
         description="""Timeout for the delete prompt operation in seconds. Defaults to 90.""",
     )
+    max_wait_time: Optional[int] = Field(
+        default=60,
+        description="""Maximum interval between polling requests in seconds. Defaults to 60.""",
+    )
 
 
 class DeletePromptConfigDict(TypedDict, total=False):
@@ -21115,6 +21119,9 @@ class DeletePromptConfigDict(TypedDict, total=False):
     timeout: Optional[int]
     """Timeout for the delete prompt operation in seconds. Defaults to 90."""
 
+    max_wait_time: Optional[int]
+    """Maximum interval between polling requests in seconds. Defaults to 60."""
+
 
 DeletePromptConfigOrDict = Union[DeletePromptConfig, DeletePromptConfigDict]
 
@@ -21262,6 +21269,14 @@ class RestoreVersionConfig(_common.BaseModel):
     http_options: Optional[genai_types.HttpOptions] = Field(
         default=None, description="""Used to override HTTP request options."""
     )
+    timeout: Optional[int] = Field(
+        default=90,
+        description="""Timeout for the restore prompt version operation in seconds. Defaults to 90.""",
+    )
+    max_wait_time: Optional[int] = Field(
+        default=60,
+        description="""Maximum interval between polling requests in seconds. Defaults to 60.""",
+    )
 
 
 class RestoreVersionConfigDict(TypedDict, total=False):
@@ -21270,6 +21285,12 @@ class RestoreVersionConfigDict(TypedDict, total=False):
     http_options: Optional[genai_types.HttpOptions]
     """Used to override HTTP request options."""
 
+    timeout: Optional[int]
+    """Timeout for the restore prompt version operation in seconds. Defaults to 90."""
+
+    max_wait_time: Optional[int]
+    """Maximum interval between polling requests in seconds. Defaults to 60."""
+
 
 RestoreVersionConfigOrDict = Union[RestoreVersionConfig, RestoreVersionConfigDict]
 
@@ -21367,6 +21388,10 @@ class UpdatePromptConfig(_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.""",
     )
+    max_wait_time: Optional[int] = Field(
+        default=60,
+        description="""The maximum interval between polling requests in seconds. If not set, the default interval is 60 seconds.""",
+    )
 
 
 class UpdatePromptConfigDict(TypedDict, total=False):
@@ -21387,6 +21412,9 @@ class UpdatePromptConfigDict(TypedDict, total=False):
     encryption_spec: Optional[genai_types.EncryptionSpec]
     """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."""
 
+    max_wait_time: Optional[int]
+    """The maximum interval between polling requests in seconds. If not set, the default interval is 60 seconds."""
+
 
 UpdatePromptConfigOrDict = Union[UpdatePromptConfig, UpdatePromptConfigDict]
 
@@ -26388,6 +26416,10 @@ class CreatePromptConfig(_common.BaseModel):
         default=None,
         description="""The display name for the prompt version. If not set, a default name with a timestamp will be used.""",
     )
+    max_wait_time: Optional[int] = Field(
+        default=60,
+        description="""The maximum interval between requests in seconds. If not set, the default interval is 60 seconds.""",
+    )
 
 
 class CreatePromptConfigDict(TypedDict, total=False):
@@ -26408,6 +26440,9 @@ class CreatePromptConfigDict(TypedDict, total=False):
     version_display_name: Optional[str]
     """The display name for the prompt version. If not set, a default name with a timestamp will be used."""
 
+    max_wait_time: Optional[int]
+    """The maximum interval between requests in seconds. If not set, the default interval is 60 seconds."""
+
 
 CreatePromptConfigOrDict = Union[CreatePromptConfig, CreatePromptConfigDict]
 
@@ -26434,6 +26469,10 @@ class CreatePromptVersionConfig(_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.""",
     )
+    max_wait_time: Optional[int] = Field(
+        default=60,
+        description="""The maximum interval between requests in seconds. If not set, the default interval is 60 seconds.""",
+    )
 
 
 class CreatePromptVersionConfigDict(TypedDict, total=False):
@@ -26454,6 +26493,9 @@ class CreatePromptVersionConfigDict(TypedDict, total=False):
     encryption_spec: Optional[genai_types.EncryptionSpec]
     """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."""
 
+    max_wait_time: Optional[int]
+    """The maximum interval between requests in seconds. If not set, the default interval is 60 seconds."""
+
 
 CreatePromptVersionConfigOrDict = Union[
     CreatePromptVersionConfig, CreatePromptVersionConfigDict