Skip to content

Commit ad63bfe

Browse files
speedstorm1copybara-github
speedstorm1
authored and
copybara-github
committed
feat: allow users to configure max wait time for prompt management queries for Python and JS
PiperOrigin-RevId: 956048771
1 parent 8b43add commit ad63bfe

2 files changed

Lines changed: 203 additions & 19 deletions

File tree

‎agentplatform/_genai/prompts.py‎

Lines changed: 161 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1505,6 +1505,9 @@ def _optimize(
15051505
self._api_client._verify_response(return_value)
15061506
return return_value
15071507

1508+
_DEFAULT_TIMEOUT = 90
1509+
_DEFAULT_MAX_WAIT_TIME = 60
1510+
15081511
def create(
15091512
self,
15101513
*,
@@ -1564,7 +1567,16 @@ def create(
15641567
)
15651568
dataset_resource_name = self._wait_for_operation(
15661569
operation=create_prompt_dataset_operation,
1567-
timeout=config.timeout if config else 90,
1570+
timeout=(
1571+
config.timeout
1572+
if config and config.timeout is not None
1573+
else self._DEFAULT_TIMEOUT
1574+
),
1575+
max_wait_time=(
1576+
config.max_wait_time
1577+
if config and config.max_wait_time is not None
1578+
else self._DEFAULT_MAX_WAIT_TIME
1579+
),
15681580
)
15691581
dataset_id = dataset_resource_name.split("/")[-1]
15701582

@@ -1636,7 +1648,16 @@ def create_version(
16361648
)
16371649
dataset_resource_name = self._wait_for_operation(
16381650
operation=create_prompt_dataset_operation,
1639-
timeout=config.timeout if config else 90,
1651+
timeout=(
1652+
config.timeout
1653+
if config and config.timeout is not None
1654+
else self._DEFAULT_TIMEOUT
1655+
),
1656+
max_wait_time=(
1657+
config.max_wait_time
1658+
if config and config.max_wait_time is not None
1659+
else self._DEFAULT_MAX_WAIT_TIME
1660+
),
16401661
)
16411662
dataset_id = dataset_resource_name.split("/")[-1]
16421663

@@ -1660,7 +1681,16 @@ def create_version(
16601681
)
16611682
dataset_version_resource_name = self._wait_for_operation(
16621683
operation=create_dataset_version_operation,
1663-
timeout=config.timeout if config else 90,
1684+
timeout=(
1685+
config.timeout
1686+
if config and config.timeout is not None
1687+
else self._DEFAULT_TIMEOUT
1688+
),
1689+
max_wait_time=(
1690+
config.max_wait_time
1691+
if config and config.max_wait_time is not None
1692+
else self._DEFAULT_MAX_WAIT_TIME
1693+
),
16641694
)
16651695

16661696
# Step 4: Get the dataset version resource and return it with the prompt
@@ -1679,12 +1709,14 @@ def _wait_for_operation(
16791709
self,
16801710
operation: types.DatasetOperation,
16811711
timeout: int,
1712+
max_wait_time: int = _DEFAULT_MAX_WAIT_TIME,
16821713
) -> str:
16831714
"""Waits for a dataset operation to complete.
16841715
16851716
Args:
16861717
operation: The dataset operation to wait for.
16871718
timeout: The maximum time to wait for the operation to complete.
1719+
max_wait_time: The maximum interval between polling requests in seconds.
16881720
16891721
Returns:
16901722
The name of the Dataset resource from the operation result.
@@ -1706,7 +1738,6 @@ def _wait_for_operation(
17061738
start_time = time.time()
17071739
sleep_duration = 5
17081740
wait_multiplier = 2
1709-
max_wait_time = 60
17101741
previous_time = time.time()
17111742

17121743
while not done:
@@ -1923,6 +1954,7 @@ def _wait_for_project_operation(
19231954
self,
19241955
operation: genai_types.ProjectOperation,
19251956
timeout: int,
1957+
max_wait_time: int = _DEFAULT_MAX_WAIT_TIME,
19261958
) -> None:
19271959
"""Waits for a dataset deletion operation to complete.
19281960
@@ -1931,6 +1963,7 @@ def _wait_for_project_operation(
19311963
Args:
19321964
operation: The project operation to wait for.
19331965
timeout: The maximum time to wait for the operation to complete.
1966+
max_wait_time: The maximum interval between polling requests in seconds.
19341967
Raises:
19351968
TimeoutError: If the operation does not complete within the timeout.
19361969
ValueError: If the operation fails.
@@ -1940,7 +1973,6 @@ def _wait_for_project_operation(
19401973
start_time = time.time()
19411974
sleep_duration = 5
19421975
wait_multiplier = 2
1943-
max_wait_time = 60
19441976
previous_time = time.time()
19451977
while not done:
19461978
if (time.time() - start_time) > timeout:
@@ -1985,7 +2017,17 @@ def delete(
19852017
config=config,
19862018
)
19872019
self._wait_for_project_operation(
1988-
operation=delete_prompt_operation, timeout=config.timeout if config else 90
2020+
operation=delete_prompt_operation,
2021+
timeout=(
2022+
config.timeout
2023+
if config and config.timeout is not None
2024+
else self._DEFAULT_TIMEOUT
2025+
),
2026+
max_wait_time=(
2027+
config.max_wait_time
2028+
if config and config.max_wait_time is not None
2029+
else self._DEFAULT_MAX_WAIT_TIME
2030+
),
19892031
)
19902032
logger.info(f"Deleted prompt with id: {prompt_id}")
19912033

@@ -2013,7 +2055,17 @@ def delete_version(
20132055
)
20142056

20152057
self._wait_for_project_operation(
2016-
operation=delete_version_operation, timeout=config.timeout if config else 90
2058+
operation=delete_version_operation,
2059+
timeout=(
2060+
config.timeout
2061+
if config and config.timeout is not None
2062+
else self._DEFAULT_TIMEOUT
2063+
),
2064+
max_wait_time=(
2065+
config.max_wait_time
2066+
if config and config.max_wait_time is not None
2067+
else self._DEFAULT_MAX_WAIT_TIME
2068+
),
20172069
)
20182070
logger.info(
20192071
f"Deleted prompt version {version_id} from prompt with id: {prompt_id}"
@@ -2040,10 +2092,20 @@ def restore_version(
20402092
restore_prompt_operation = self._restore_version(
20412093
dataset_id=prompt_id,
20422094
version_id=version_id,
2095+
config=config,
20432096
)
20442097
self._wait_for_project_operation(
20452098
operation=restore_prompt_operation,
2046-
timeout=90,
2099+
timeout=(
2100+
config.timeout
2101+
if config and config.timeout is not None
2102+
else self._DEFAULT_TIMEOUT
2103+
),
2104+
max_wait_time=(
2105+
config.max_wait_time
2106+
if config and config.max_wait_time is not None
2107+
else self._DEFAULT_MAX_WAIT_TIME
2108+
),
20472109
)
20482110
dataset_version_resource = self._get_dataset_version_resource(
20492111
dataset_id=prompt_id,
@@ -2070,7 +2132,7 @@ def _wait_for_completion(self, job_name: str) -> types.CustomJob:
20702132

20712133
log_wait = 5
20722134
wait_multiplier = 2
2073-
max_wait_time = 60
2135+
max_wait_time = self._DEFAULT_MAX_WAIT_TIME
20742136
previous_time = time.time()
20752137

20762138
job = self._get_custom_job(name=job_name)
@@ -2400,7 +2462,16 @@ def update(
24002462
)
24012463
dataset_version_resource_name = self._wait_for_operation(
24022464
operation=create_dataset_version_operation,
2403-
timeout=config.timeout if config else 90,
2465+
timeout=(
2466+
config.timeout
2467+
if config and config.timeout is not None
2468+
else self._DEFAULT_TIMEOUT
2469+
),
2470+
max_wait_time=(
2471+
config.max_wait_time
2472+
if config and config.max_wait_time is not None
2473+
else self._DEFAULT_MAX_WAIT_TIME
2474+
),
24042475
)
24052476
dataset_version_id = dataset_version_resource_name.split("/")[-1]
24062477

@@ -3499,6 +3570,9 @@ async def _optimize(
34993570
self._api_client._verify_response(return_value)
35003571
return return_value
35013572

3573+
_DEFAULT_TIMEOUT = 90
3574+
_DEFAULT_MAX_WAIT_TIME = 60
3575+
35023576
async def create(
35033577
self,
35043578
*,
@@ -3558,7 +3632,16 @@ async def create(
35583632
)
35593633
dataset_resource_name = await self._wait_for_operation(
35603634
operation=create_prompt_dataset_operation,
3561-
timeout=config.timeout if config else 90,
3635+
timeout=(
3636+
config.timeout
3637+
if config and config.timeout is not None
3638+
else self._DEFAULT_TIMEOUT
3639+
),
3640+
max_wait_time=(
3641+
config.max_wait_time
3642+
if config and config.max_wait_time is not None
3643+
else self._DEFAULT_MAX_WAIT_TIME
3644+
),
35623645
)
35633646
dataset_id = dataset_resource_name.split("/")[-1]
35643647

@@ -3629,7 +3712,16 @@ async def create_version(
36293712
)
36303713
dataset_resource_name = await self._wait_for_operation(
36313714
operation=create_prompt_dataset_operation,
3632-
timeout=config.timeout if config else 90,
3715+
timeout=(
3716+
config.timeout
3717+
if config and config.timeout is not None
3718+
else self._DEFAULT_TIMEOUT
3719+
),
3720+
max_wait_time=(
3721+
config.max_wait_time
3722+
if config and config.max_wait_time is not None
3723+
else self._DEFAULT_MAX_WAIT_TIME
3724+
),
36333725
)
36343726
dataset_id = dataset_resource_name.split("/")[-1]
36353727

@@ -3653,7 +3745,16 @@ async def create_version(
36533745
)
36543746
dataset_version_resource_name = await self._wait_for_operation(
36553747
operation=create_dataset_version_operation,
3656-
timeout=config.timeout if config else 90,
3748+
timeout=(
3749+
config.timeout
3750+
if config and config.timeout is not None
3751+
else self._DEFAULT_TIMEOUT
3752+
),
3753+
max_wait_time=(
3754+
config.max_wait_time
3755+
if config and config.max_wait_time is not None
3756+
else self._DEFAULT_MAX_WAIT_TIME
3757+
),
36573758
)
36583759

36593760
# Step 4: Get the dataset version resource and return it with the prompt
@@ -3740,7 +3841,16 @@ async def update(
37403841
)
37413842
dataset_version_resource_name = await self._wait_for_operation(
37423843
operation=create_dataset_version_operation,
3743-
timeout=config.timeout if config else 90,
3844+
timeout=(
3845+
config.timeout
3846+
if config and config.timeout is not None
3847+
else self._DEFAULT_TIMEOUT
3848+
),
3849+
max_wait_time=(
3850+
config.max_wait_time
3851+
if config and config.max_wait_time is not None
3852+
else self._DEFAULT_MAX_WAIT_TIME
3853+
),
37443854
)
37453855
dataset_version_id = dataset_version_resource_name.split("/")[-1]
37463856

@@ -3760,12 +3870,14 @@ async def _wait_for_operation(
37603870
self,
37613871
operation: types.DatasetOperation,
37623872
timeout: int,
3873+
max_wait_time: int = _DEFAULT_MAX_WAIT_TIME,
37633874
) -> str:
37643875
"""Waits for a dataset operation to complete.
37653876
37663877
Args:
37673878
operation: The dataset operation to wait for.
37683879
timeout: The maximum time to wait for the operation to complete.
3880+
max_wait_time: The maximum interval between polling requests in seconds.
37693881
37703882
Returns:
37713883
The name of the Dataset resource from the operation result.
@@ -3787,7 +3899,6 @@ async def _wait_for_operation(
37873899
start_time = time.time()
37883900
sleep_duration = 5
37893901
wait_multiplier = 2
3790-
max_wait_time = 60
37913902
previous_time = time.time()
37923903

37933904
while not done:
@@ -3885,6 +3996,7 @@ async def _wait_for_project_operation(
38853996
self,
38863997
operation: genai_types.ProjectOperation,
38873998
timeout: int,
3999+
max_wait_time: int = _DEFAULT_MAX_WAIT_TIME,
38884000
) -> None:
38894001
"""Waits for a dataset deletion operation to complete.
38904002
@@ -3893,6 +4005,7 @@ async def _wait_for_project_operation(
38934005
Args:
38944006
operation: The project operation to wait for.
38954007
timeout: The maximum time to wait for the operation to complete.
4008+
max_wait_time: The maximum interval between polling requests in seconds.
38964009
Raises:
38974010
TimeoutError: If the operation does not complete within the timeout.
38984011
ValueError: If the operation fails.
@@ -3902,7 +4015,6 @@ async def _wait_for_project_operation(
39024015
start_time = time.time()
39034016
sleep_duration = 5
39044017
wait_multiplier = 2
3905-
max_wait_time = 60
39064018
previous_time = time.time()
39074019
while not done:
39084020
if (time.time() - start_time) > timeout:
@@ -3947,7 +4059,17 @@ async def delete(
39474059
config=config,
39484060
)
39494061
await self._wait_for_project_operation(
3950-
operation=delete_prompt_operation, timeout=config.timeout if config else 90
4062+
operation=delete_prompt_operation,
4063+
timeout=(
4064+
config.timeout
4065+
if config and config.timeout is not None
4066+
else self._DEFAULT_TIMEOUT
4067+
),
4068+
max_wait_time=(
4069+
config.max_wait_time
4070+
if config and config.max_wait_time is not None
4071+
else self._DEFAULT_MAX_WAIT_TIME
4072+
),
39514073
)
39524074
logger.info(f"Deleted prompt with id: {prompt_id}")
39534075

@@ -3975,7 +4097,17 @@ async def delete_version(
39754097
)
39764098

39774099
await self._wait_for_project_operation(
3978-
operation=delete_version_operation, timeout=config.timeout if config else 90
4100+
operation=delete_version_operation,
4101+
timeout=(
4102+
config.timeout
4103+
if config and config.timeout is not None
4104+
else self._DEFAULT_TIMEOUT
4105+
),
4106+
max_wait_time=(
4107+
config.max_wait_time
4108+
if config and config.max_wait_time is not None
4109+
else self._DEFAULT_MAX_WAIT_TIME
4110+
),
39794111
)
39804112
logger.info(
39814113
f"Deleted prompt version {version_id} from prompt with id: {prompt_id}"
@@ -4111,10 +4243,20 @@ async def restore_version(
41114243
restore_prompt_operation = await self._restore_version(
41124244
dataset_id=prompt_id,
41134245
version_id=version_id,
4246+
config=config,
41144247
)
41154248
await self._wait_for_project_operation(
41164249
operation=restore_prompt_operation,
4117-
timeout=90,
4250+
timeout=(
4251+
config.timeout
4252+
if config and config.timeout is not None
4253+
else self._DEFAULT_TIMEOUT
4254+
),
4255+
max_wait_time=(
4256+
config.max_wait_time
4257+
if config and config.max_wait_time is not None
4258+
else self._DEFAULT_MAX_WAIT_TIME
4259+
),
41184260
)
41194261
dataset_version_resource = await self._get_dataset_version_resource(
41204262
dataset_id=prompt_id,

0 commit comments

Comments
 (0)