From ad63bfee932843ad8a0477ebb07b0159269c74bb Mon Sep 17 00:00:00 2001 From: Ayush Agrawal Date: Wed, 29 Jul 2026 13:13:00 -0700 Subject: [PATCH] feat: allow users to configure max wait time for prompt management queries for Python and JS PiperOrigin-RevId: 956048771 --- agentplatform/_genai/prompts.py | 180 ++++++++++++++++++++++++--- agentplatform/_genai/types/common.py | 42 +++++++ 2 files changed, 203 insertions(+), 19 deletions(-) diff --git a/agentplatform/_genai/prompts.py b/agentplatform/_genai/prompts.py index 86bd483728..43e1ac1944 100644 --- a/agentplatform/_genai/prompts.py +++ b/agentplatform/_genai/prompts.py @@ -1505,6 +1505,9 @@ def _optimize( self._api_client._verify_response(return_value) return return_value + _DEFAULT_TIMEOUT = 90 + _DEFAULT_MAX_WAIT_TIME = 60 + def create( self, *, @@ -1564,7 +1567,16 @@ 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 self._DEFAULT_TIMEOUT + ), + max_wait_time=( + config.max_wait_time + if config and config.max_wait_time is not None + else self._DEFAULT_MAX_WAIT_TIME + ), ) dataset_id = dataset_resource_name.split("/")[-1] @@ -1636,7 +1648,16 @@ 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 self._DEFAULT_TIMEOUT + ), + max_wait_time=( + config.max_wait_time + if config and config.max_wait_time is not None + else self._DEFAULT_MAX_WAIT_TIME + ), ) dataset_id = dataset_resource_name.split("/")[-1] @@ -1660,7 +1681,16 @@ 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 self._DEFAULT_TIMEOUT + ), + max_wait_time=( + config.max_wait_time + if config and config.max_wait_time is not None + else self._DEFAULT_MAX_WAIT_TIME + ), ) # Step 4: Get the dataset version resource and return it with the prompt @@ -1679,12 +1709,14 @@ def _wait_for_operation( self, operation: types.DatasetOperation, timeout: int, + max_wait_time: int = _DEFAULT_MAX_WAIT_TIME, ) -> 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 +1738,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 +1954,7 @@ def _wait_for_project_operation( self, operation: genai_types.ProjectOperation, timeout: int, + max_wait_time: int = _DEFAULT_MAX_WAIT_TIME, ) -> None: """Waits for a dataset deletion operation to complete. @@ -1931,6 +1963,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 +1973,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 +2017,17 @@ 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 self._DEFAULT_TIMEOUT + ), + max_wait_time=( + config.max_wait_time + if config and config.max_wait_time is not None + else self._DEFAULT_MAX_WAIT_TIME + ), ) logger.info(f"Deleted prompt with id: {prompt_id}") @@ -2013,7 +2055,17 @@ 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 self._DEFAULT_TIMEOUT + ), + max_wait_time=( + config.max_wait_time + if config and config.max_wait_time is not None + else self._DEFAULT_MAX_WAIT_TIME + ), ) logger.info( f"Deleted prompt version {version_id} from prompt with id: {prompt_id}" @@ -2040,10 +2092,20 @@ 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 self._DEFAULT_TIMEOUT + ), + max_wait_time=( + config.max_wait_time + if config and config.max_wait_time is not None + else self._DEFAULT_MAX_WAIT_TIME + ), ) dataset_version_resource = self._get_dataset_version_resource( dataset_id=prompt_id, @@ -2070,7 +2132,7 @@ def _wait_for_completion(self, job_name: str) -> types.CustomJob: log_wait = 5 wait_multiplier = 2 - max_wait_time = 60 + max_wait_time = self._DEFAULT_MAX_WAIT_TIME previous_time = time.time() job = self._get_custom_job(name=job_name) @@ -2400,7 +2462,16 @@ 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 self._DEFAULT_TIMEOUT + ), + max_wait_time=( + config.max_wait_time + if config and config.max_wait_time is not None + else self._DEFAULT_MAX_WAIT_TIME + ), ) dataset_version_id = dataset_version_resource_name.split("/")[-1] @@ -3499,6 +3570,9 @@ async def _optimize( self._api_client._verify_response(return_value) return return_value + _DEFAULT_TIMEOUT = 90 + _DEFAULT_MAX_WAIT_TIME = 60 + async def create( self, *, @@ -3558,7 +3632,16 @@ 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 self._DEFAULT_TIMEOUT + ), + max_wait_time=( + config.max_wait_time + if config and config.max_wait_time is not None + else self._DEFAULT_MAX_WAIT_TIME + ), ) dataset_id = dataset_resource_name.split("/")[-1] @@ -3629,7 +3712,16 @@ 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 self._DEFAULT_TIMEOUT + ), + max_wait_time=( + config.max_wait_time + if config and config.max_wait_time is not None + else self._DEFAULT_MAX_WAIT_TIME + ), ) dataset_id = dataset_resource_name.split("/")[-1] @@ -3653,7 +3745,16 @@ 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 self._DEFAULT_TIMEOUT + ), + max_wait_time=( + config.max_wait_time + if config and config.max_wait_time is not None + else self._DEFAULT_MAX_WAIT_TIME + ), ) # Step 4: Get the dataset version resource and return it with the prompt @@ -3740,7 +3841,16 @@ 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 self._DEFAULT_TIMEOUT + ), + max_wait_time=( + config.max_wait_time + if config and config.max_wait_time is not None + else self._DEFAULT_MAX_WAIT_TIME + ), ) dataset_version_id = dataset_version_resource_name.split("/")[-1] @@ -3760,12 +3870,14 @@ async def _wait_for_operation( self, operation: types.DatasetOperation, timeout: int, + max_wait_time: int = _DEFAULT_MAX_WAIT_TIME, ) -> 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 +3899,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 +3996,7 @@ async def _wait_for_project_operation( self, operation: genai_types.ProjectOperation, timeout: int, + max_wait_time: int = _DEFAULT_MAX_WAIT_TIME, ) -> None: """Waits for a dataset deletion operation to complete. @@ -3893,6 +4005,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 +4015,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 +4059,17 @@ 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 self._DEFAULT_TIMEOUT + ), + max_wait_time=( + config.max_wait_time + if config and config.max_wait_time is not None + else self._DEFAULT_MAX_WAIT_TIME + ), ) logger.info(f"Deleted prompt with id: {prompt_id}") @@ -3975,7 +4097,17 @@ 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 self._DEFAULT_TIMEOUT + ), + max_wait_time=( + config.max_wait_time + if config and config.max_wait_time is not None + else self._DEFAULT_MAX_WAIT_TIME + ), ) logger.info( f"Deleted prompt version {version_id} from prompt with id: {prompt_id}" @@ -4111,10 +4243,20 @@ 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 self._DEFAULT_TIMEOUT + ), + max_wait_time=( + config.max_wait_time + if config and config.max_wait_time is not None + else self._DEFAULT_MAX_WAIT_TIME + ), ) 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 e2ff847bac..24265d7dab 100644 --- a/agentplatform/_genai/types/common.py +++ b/agentplatform/_genai/types/common.py @@ -21554,6 +21554,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): @@ -21565,6 +21569,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] @@ -21712,6 +21719,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): @@ -21720,6 +21735,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] @@ -21817,6 +21838,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): @@ -21837,6 +21862,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] @@ -26958,6 +26986,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): @@ -26978,6 +27010,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] @@ -27004,6 +27039,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): @@ -27024,6 +27063,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