diff --git a/tests/unit/vertexai/genai/replays/test_create_agent_engine_sandbox.py b/tests/unit/vertexai/genai/replays/test_create_agent_engine_sandbox.py index 39124d9775..70c3e73269 100644 --- a/tests/unit/vertexai/genai/replays/test_create_agent_engine_sandbox.py +++ b/tests/unit/vertexai/genai/replays/test_create_agent_engine_sandbox.py @@ -25,6 +25,7 @@ def test_create_sandbox(client): operation = client.agent_engines.sandboxes.create( name=agent_engine.api_resource.name, + poll_interval_seconds=1, spec={ "code_execution_environment": { "machineConfig": "MACHINE_CONFIG_VCPU4_RAM4GIB" diff --git a/vertexai/_genai/sandboxes.py b/vertexai/_genai/sandboxes.py index c88be4e269..88eb1d6b6f 100644 --- a/vertexai/_genai/sandboxes.py +++ b/vertexai/_genai/sandboxes.py @@ -540,6 +540,7 @@ def create( self, *, name: str, + poll_interval_seconds: float = 0.1, spec: Optional[types.SandboxEnvironmentSpecOrDict] = None, config: Optional[types.CreateAgentEngineSandboxConfigOrDict] = None, ) -> types.AgentEngineSandboxOperation: @@ -582,7 +583,7 @@ def create( operation = _agent_engines_utils._await_operation( operation_name=operation.name, get_operation_fn=self._get_sandbox_operation, - poll_interval_seconds=0.1, + poll_interval_seconds=poll_interval_seconds, ) # We need to make a call to get the sandbox because the operation # response might not contain the relevant fields.