Skip to content

Commit a28bd14

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
feat: Add support for wait for operation complete for sandbox sub-modules
PiperOrigin-RevId: 908889456
1 parent b8eaefb commit a28bd14

10 files changed

Lines changed: 282 additions & 45 deletions

‎tests/unit/vertexai/genai/replays/test_ae_sandbox_snapshots_create.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121

2222

2323
def test_create_sandbox_snapshot(client):
24-
snapshot = client.agent_engines.sandboxes.snapshots.create(
24+
snapshot = client.agent_engines.sandboxes.snapshots._create(
2525
source_sandbox_environment_name="projects/802583348448/locations/us-central1/reasoningEngines/6130241318758121472/sandboxEnvironments/525190525100228608",
2626
config={
2727
"display_name": "test_snapshot",
@@ -36,5 +36,5 @@ def test_create_sandbox_snapshot(client):
3636
pytestmark = pytest_helper.setup(
3737
file=__file__,
3838
globals_for_file=globals(),
39-
test_method="agent_engines.sandboxes.snapshots.create",
39+
test_method="agent_engines.sandboxes.snapshots._create",
4040
)

‎tests/unit/vertexai/genai/replays/test_ae_sandbox_snapshots_delete.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121

2222

2323
def test_delete_sandbox_snapshot(client):
24-
result = client.agent_engines.sandboxes.snapshots.delete(
24+
result = client.agent_engines.sandboxes.snapshots._delete(
2525
name="projects/802583348448/locations/us-central1/reasoningEngines/6130241318758121472/sandboxEnvironmentSnapshots/421086565159141376",
2626
)
2727

@@ -31,5 +31,5 @@ def test_delete_sandbox_snapshot(client):
3131
pytestmark = pytest_helper.setup(
3232
file=__file__,
3333
globals_for_file=globals(),
34-
test_method="agent_engines.sandboxes.snapshots.delete",
34+
test_method="agent_engines.sandboxes.snapshots._delete",
3535
)

‎tests/unit/vertexai/genai/replays/test_ae_sandbox_snapshots_get.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,8 @@
2121

2222

2323
def test_get_sandbox_snapshot(client):
24-
snapshot_name = "projects/802583348448/locations/us-central1/reasoningEngines/6130241318758121472/sandboxEnvironmentSnapshots/421086565159141376"
25-
fetched_snapshot = client.agent_engines.sandboxes.snapshots.get(
24+
snapshot_name = "projects/802583348448/locations/us-central1/reasoningEngines/6130241318758121472/sandboxEnvironmentSnapshots/2433069698686910464"
25+
fetched_snapshot = client.agent_engines.sandboxes.snapshots._get(
2626
name=snapshot_name,
2727
)
2828

@@ -33,5 +33,5 @@ def test_get_sandbox_snapshot(client):
3333
pytestmark = pytest_helper.setup(
3434
file=__file__,
3535
globals_for_file=globals(),
36-
test_method="agent_engines.sandboxes.snapshots.get",
36+
test_method="agent_engines.sandboxes.snapshots._get",
3737
)

‎tests/unit/vertexai/genai/replays/test_ae_sandbox_snapshots_list.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121

2222

2323
def test_list_sandbox_snapshots(client):
24-
snapshots_list_operation = client.agent_engines.sandboxes.snapshots.list(
24+
snapshots_list_operation = client.agent_engines.sandboxes.snapshots._list(
2525
name="projects/802583348448/locations/us-central1/reasoningEngines/6130241318758121472",
2626
)
2727

@@ -37,5 +37,5 @@ def test_list_sandbox_snapshots(client):
3737
pytestmark = pytest_helper.setup(
3838
file=__file__,
3939
globals_for_file=globals(),
40-
test_method="agent_engines.sandboxes.snapshots.list",
40+
test_method="agent_engines.sandboxes.snapshots._list",
4141
)

‎tests/unit/vertexai/genai/replays/test_ae_sandbox_templates_byoc_create.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ def test_sandbox_templates_byoc_create(client):
5050
}
5151
sandbox_template_operation = client.agent_engines.sandboxes.templates.create(
5252
name=(
53-
"projects/254005681254/locations/us-central1/reasoningEngines/208148546254274560"
53+
"projects/802583348448/locations/us-central1/reasoningEngines/6130241318758121472"
5454
),
5555
display_name="Test Sandbox Template 1",
5656
config=config,

‎tests/unit/vertexai/genai/replays/test_ae_sandbox_templates_default_create.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -27,9 +27,9 @@ def test_sandbox_templates_default_create(client):
2727
"internet_access": True,
2828
},
2929
}
30-
sandbox_template_operation = client.agent_engines.sandboxes.templates.create(
30+
sandbox_template_operation = client.agent_engines.sandboxes.templates._create(
3131
name=(
32-
"projects/254005681254/locations/us-central1/reasoningEngines/208148546254274560"
32+
"projects/802583348448/locations/us-central1/reasoningEngines/6130241318758121472"
3333
),
3434
display_name="Test Sandbox Template 1",
3535
config=config,
@@ -42,5 +42,5 @@ def test_sandbox_templates_default_create(client):
4242
pytestmark = pytest_helper.setup(
4343
file=__file__,
4444
globals_for_file=globals(),
45-
test_method="agent_engines.sandboxes.templates.create",
45+
test_method="agent_engines.sandboxes.templates._create",
4646
)

‎tests/unit/vertexai/genai/replays/test_ae_sandbox_templates_list.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919

2020

2121
def test_sandbox_templates_list(client):
22-
sandbox_templates_list_operation = client.agent_engines.sandboxes.templates.list(
22+
sandbox_templates_list_operation = client.agent_engines.sandboxes.templates._list(
2323
name=(
2424
"projects/254005681254/locations/us-central1/reasoningEngines/208148546254274560"
2525
),
@@ -36,5 +36,5 @@ def test_sandbox_templates_list(client):
3636
pytestmark = pytest_helper.setup(
3737
file=__file__,
3838
globals_for_file=globals(),
39-
test_method="agent_engines.sandboxes.templates.list",
39+
test_method="agent_engines.sandboxes.templates._list",
4040
)

‎vertexai/_genai/sandbox_snapshots.py‎

Lines changed: 133 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -15,20 +15,25 @@
1515

1616
# Code generated by the Google Gen AI SDK generator DO NOT EDIT.
1717

18+
import functools
1819
import json
1920
import logging
20-
from typing import Any, Optional, Union
21+
from typing import Any, Iterator, Optional, Union
2122
from urllib.parse import urlencode
2223

2324
from google.genai import _api_module
2425
from google.genai import _common
2526
from google.genai._common import get_value_by_path as getv
2627
from google.genai._common import set_value_by_path as setv
28+
from google.genai.pagers import Pager
2729

30+
from . import _agent_engines_utils
2831
from . import types
2932

3033
logger = logging.getLogger("vertexai_genai.sandboxsnapshots")
3134

35+
logger.setLevel(logging.INFO)
36+
3237

3338
def _CreateAgentEngineSandboxSnapshotConfig_to_vertex(
3439
from_object: Union[dict[str, Any], object],
@@ -144,14 +149,21 @@ def _ListSandboxEnvironmentSnapshotsRequestParameters_to_vertex(
144149
class SandboxSnapshots(_api_module.BaseModule):
145150
"""Sandbox environment snapshot commands."""
146151

147-
def create(
152+
def _create(
148153
self,
149154
*,
150155
source_sandbox_environment_name: str,
151156
config: Optional[types.CreateAgentEngineSandboxSnapshotConfigOrDict] = None,
152157
) -> types.AgentEngineSandboxSnapshotOperation:
153158
"""
154-
Snapshots a sandbox environment.
159+
Snapshots an existing sandbox environment.
160+
161+
Args:
162+
source_sandbox_environment_name (str):
163+
Required. The name of the sandbox environment to snapshot.
164+
projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sandboxEnvironments/{sandbox_environment_id}
165+
config (CreateAgentEngineSandboxSnapshotConfig):
166+
Optional. The configuration for the sandbox snapshot.
155167
156168
"""
157169

@@ -219,7 +231,7 @@ def create(
219231
self._api_client._verify_response(return_value)
220232
return return_value
221233

222-
def delete(
234+
def _delete(
223235
self,
224236
*,
225237
name: str,
@@ -294,7 +306,7 @@ def delete(
294306
self._api_client._verify_response(return_value)
295307
return return_value
296308

297-
def get(
309+
def _get(
298310
self,
299311
*,
300312
name: str,
@@ -369,7 +381,7 @@ def get(
369381
self._api_client._verify_response(return_value)
370382
return return_value
371383

372-
def list(
384+
def _list(
373385
self,
374386
*,
375387
name: str,
@@ -444,7 +456,7 @@ def list(
444456
self._api_client._verify_response(return_value)
445457
return return_value
446458

447-
def _get_sandbox_snapshot_operation(
459+
def get_sandbox_snapshot_operation(
448460
self,
449461
*,
450462
operation_name: str,
@@ -514,18 +526,126 @@ def _get_sandbox_snapshot_operation(
514526
self._api_client._verify_response(return_value)
515527
return return_value
516528

529+
def create(
530+
self,
531+
*,
532+
source_sandbox_environment_name: str,
533+
config: Optional[types.CreateAgentEngineSandboxSnapshotConfigOrDict] = None,
534+
poll_interval_seconds: float = 0.1,
535+
) -> types.AgentEngineSandboxSnapshotOperation:
536+
"""Snapshots an existing sandbox environment.
537+
538+
Args:
539+
source_sandbox_environment_name (str):
540+
Required. The name of the sandbox environment to snapshot.
541+
projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sandboxEnvironments/{sandbox_environment_id}
542+
config (CreateAgentEngineSandboxSnapshotConfig):
543+
Optional. The configuration for the sandbox snapshot.
544+
poll_interval_seconds (int):
545+
Optional. Seconds to wait between polling for operation status. Defaults to 0.1.
546+
547+
Returns:
548+
AgentEngineSandboxSnapshotOperation: The operation for creating the sandbox snapshot.
549+
"""
550+
operation = self._create(
551+
source_sandbox_environment_name=source_sandbox_environment_name,
552+
config=config,
553+
)
554+
if config is None:
555+
config = types.CreateAgentEngineSandboxSnapshotConfig()
556+
elif isinstance(config, dict):
557+
config = types.CreateAgentEngineSandboxSnapshotConfig.model_validate(config)
558+
if config.wait_for_completion:
559+
if not operation.done:
560+
operation = _agent_engines_utils._await_operation(
561+
operation_name=operation.name,
562+
get_operation_fn=self.get_sandbox_snapshot_operation,
563+
poll_interval_seconds=poll_interval_seconds,
564+
)
565+
# We need to make a call to get the sandbox snapshot because the operation
566+
# response might not contain the relevant fields.
567+
if not operation.response:
568+
raise ValueError("Error retrieving sandbox snapshot.")
569+
operation.response = self.get(name=operation.response.name)
570+
return operation
571+
572+
def list(
573+
self,
574+
*,
575+
name: str,
576+
config: Optional[types.ListSandboxEnvironmentSnapshotsConfigOrDict] = None,
577+
) -> Iterator[types.SandboxEnvironmentSnapshot]:
578+
"""Lists Agent Engine sandbox snapshots.
579+
580+
Args:
581+
name (str):
582+
Required. The name of the agent engine to list sandbox snapshots for.
583+
projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sandboxEnvironmentSnapshots/{sandbox_snapshot_id}
584+
config (ListSandboxEnvironmentSnapshotsConfig):
585+
Optional. The configuration for the sandbox snapshots to list.
586+
587+
Returns:
588+
Iterable[SandboxEnvironmentSnapshot]: An iterable of agent engine sandbox snapshots.
589+
"""
590+
return Pager(
591+
"sandbox_environment_snapshots",
592+
functools.partial(self._list, name=name),
593+
self._list(name=name, config=config),
594+
config,
595+
)
596+
597+
def get(
598+
self,
599+
*,
600+
name: str,
601+
config: Optional[types.GetSandboxEnvironmentSnapshotConfigOrDict] = None,
602+
) -> types.SandboxEnvironmentSnapshot:
603+
"""Gets a sandbox snapshot in the Agent Engine.
604+
Args:
605+
name (str):
606+
Required. A fully-qualified resource name or ID such as
607+
projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sandboxEnvironmentSnapshots/{snapshot_id}
608+
or a shortened name such as "reasoningEngines/{resource_id}/sandboxEnvironmentSnapshots/{snapshot_id}".
609+
config (GetSandboxEnvironmentSnapshotConfigOrDict):
610+
Optional. The configuration for the sandbox snapshot to get.
611+
"""
612+
return self._get(name=name, config=config)
613+
614+
def delete(
615+
self,
616+
*,
617+
name: str,
618+
config: Optional[types.DeleteSandboxEnvironmentSnapshotConfigOrDict] = None,
619+
) -> types.DeleteSandboxEnvironmentSnapshotOperation:
620+
"""Deletes a sandbox snapshot in the Agent Engine.
621+
Args:
622+
name (str):
623+
Required. The name of the sandbox snapshot to delete.
624+
Format: projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sandboxEnvironmentSnapshots/{snapshot_id}
625+
config (DeleteSandboxEnvironmentSnapshotConfigOrDict):
626+
Optional. Configuration for the delete operation.
627+
"""
628+
return self._delete(name=name, config=config)
629+
517630

518631
class AsyncSandboxSnapshots(_api_module.BaseModule):
519632
"""Sandbox environment snapshot commands."""
520633

521-
async def create(
634+
async def _create(
522635
self,
523636
*,
524637
source_sandbox_environment_name: str,
525638
config: Optional[types.CreateAgentEngineSandboxSnapshotConfigOrDict] = None,
526639
) -> types.AgentEngineSandboxSnapshotOperation:
527640
"""
528-
Snapshots a sandbox environment.
641+
Snapshots an existing sandbox environment.
642+
643+
Args:
644+
source_sandbox_environment_name (str):
645+
Required. The name of the sandbox environment to snapshot.
646+
projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sandboxEnvironments/{sandbox_environment_id}
647+
config (CreateAgentEngineSandboxSnapshotConfig):
648+
Optional. The configuration for the sandbox snapshot.
529649
530650
"""
531651

@@ -595,7 +715,7 @@ async def create(
595715
self._api_client._verify_response(return_value)
596716
return return_value
597717

598-
async def delete(
718+
async def _delete(
599719
self,
600720
*,
601721
name: str,
@@ -672,7 +792,7 @@ async def delete(
672792
self._api_client._verify_response(return_value)
673793
return return_value
674794

675-
async def get(
795+
async def _get(
676796
self,
677797
*,
678798
name: str,
@@ -749,7 +869,7 @@ async def get(
749869
self._api_client._verify_response(return_value)
750870
return return_value
751871

752-
async def list(
872+
async def _list(
753873
self,
754874
*,
755875
name: str,
@@ -826,7 +946,7 @@ async def list(
826946
self._api_client._verify_response(return_value)
827947
return return_value
828948

829-
async def _get_sandbox_snapshot_operation(
949+
async def get_sandbox_snapshot_operation(
830950
self,
831951
*,
832952
operation_name: str,

0 commit comments

Comments
 (0)