From f85b9f9615d079af97b08bd2048d9c1557666ba9 Mon Sep 17 00:00:00 2001 From: Ayush Agrawal Date: Tue, 7 Jul 2026 19:56:03 -0700 Subject: [PATCH] feat: allow users to configure max wait time for prompt management queries for Python and JS FUTURE_COPYBARA_INTEGRATE_REVIEW=https://github.com/googleapis/python-aiplatform/pull/7003 from googleapis:release-please--branches--main 5905f2d33aa57c08026c20aa59e1a6c48b2afa05 PiperOrigin-RevId: 944229591 --- agentplatform/_genai/prompts.py | 88 ++++++++++++++++++++++------ agentplatform/_genai/types/common.py | 42 +++++++++++++ 2 files changed, 112 insertions(+), 18 deletions(-) diff --git a/agentplatform/_genai/prompts.py b/agentplatform/_genai/prompts.py index 86bd483728..ea9173be95 100644 --- a/agentplatform/_genai/prompts.py +++ b/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, diff --git a/agentplatform/_genai/types/common.py b/agentplatform/_genai/types/common.py index dd57c62b1e..5c9fe8c6b2 100644 --- a/agentplatform/_genai/types/common.py +++ b/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