Skip to content

feat: Add a top-level client.sandboxes accessor to the Vertex AI GenAI SDK for Java, with environments, templates, and snapshots submodules. #147

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: Add a top-level client.sandboxes accessor to the Vertex AI GenAI SDK for Java, with environments, templates, and snapshots submodules. #147

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
195 changes: 103 additions & 92 deletions ...java/com/google/cloud/agentplatform/examples/AgentEngineSandboxTemplatesAndSnapshots.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,193 +36,204 @@
package com.google.cloud.agentplatform.examples;

import com.google.cloud.agentplatform.Client;
import com.google.cloud.agentplatform.types.AgentEngineOperation;
import com.google.cloud.agentplatform.types.AgentEngineSandboxOperation;
import com.google.cloud.agentplatform.types.AgentEngineSandboxSnapshotOperation;
import com.google.cloud.agentplatform.types.CreateAgentEngineConfig;
import com.google.cloud.agentplatform.types.CreateAgentEngineSandboxConfig;
import com.google.cloud.agentplatform.types.CreateAgentEngineSandboxSnapshotConfig;
import com.google.cloud.agentplatform.types.CreateSandboxEnvironmentTemplateConfig;
import com.google.cloud.agentplatform.types.DeleteAgentEngineOperation;
import com.google.cloud.agentplatform.types.DeleteAgentEngineSandboxOperation;
import com.google.cloud.agentplatform.types.DeleteSandboxEnvironmentSnapshotOperation;
import com.google.cloud.agentplatform.types.DeleteSandboxEnvironmentTemplateOperation;
import com.google.cloud.agentplatform.types.ListSandboxEnvironmentSnapshotsResponse;
import com.google.cloud.agentplatform.types.ListSandboxEnvironmentTemplatesResponse;
import com.google.cloud.agentplatform.types.SandboxEnvironmentSnapshot;
import com.google.cloud.agentplatform.types.SandboxEnvironmentSpec;
import com.google.cloud.agentplatform.types.SandboxEnvironmentSpecCodeExecutionEnvironment;
import com.google.cloud.agentplatform.types.SandboxEnvironmentTemplate;
import com.google.cloud.agentplatform.types.SandboxEnvironmentTemplateDefaultContainerEnvironment;
import com.google.cloud.agentplatform.types.SandboxEnvironmentTemplateEgressControlConfig;
import com.google.cloud.agentplatform.types.SandboxEnvironmentTemplateOperation;
import com.google.common.collect.ImmutableMap;
import java.time.Duration;
import java.util.List;

/**
* An example of using the Java SDK to perform operations on Agent Engine sandbox environment
* templates and snapshots.
* An example of using the Java SDK to perform operations on Agent Platform sandbox environments,
* templates, and snapshots, using only the public API surface.
*
* <p>The sandbox API is reached from the top-level {@code client.sandboxes} accessor, with its
* {@code environments}, {@code templates}, and {@code snapshots} submodules. The parent agent
* engine name is omitted, so the SDK lazily reuses a shared default agent engine for this project +
* location.
*
* <p>The flow mirrors the canonical usage: create a computer-use sandbox template, create a sandbox
* environment from that template, snapshot the running sandbox, then restore a new sandbox from the
* snapshot. Snapshots can only be taken of template-based (e.g. computer-use) sandboxes.
*/
public final class AgentEngineSandboxTemplatesAndSnapshots {
public static void main(String[] args) {
// Instantiate the client. It reads from the environment variables `GOOGLE_CLOUD_LOCATION` and
// `GOOGLE_CLOUD_PROJECT`.
Client client = new Client();

System.out.println("Creating a temporary Agent Engine...");
CreateAgentEngineConfig config =
CreateAgentEngineConfig.builder()
.labels(ImmutableMap.of("test-label", "test-value"))
.build();
AgentEngineOperation createOperation = client.agentEngines.privateCreate(config);
while (!createOperation.done().filter(Boolean::booleanValue).isPresent()) {
sleep(10000);
createOperation =
client.agentEngines.privateGetAgentOperation(createOperation.name().get(), null);
}
String agentEngineName = createOperation.response().get().name().get();
System.out.println("Created Agent Engine: " + agentEngineName);

try {
runSandboxTemplatesExample(client, agentEngineName);
runSandboxSnapshotsExample(client, agentEngineName);
} finally {
System.out.println("\nCleaning up Agent Engine...");
DeleteAgentEngineOperation deleteOperation =
client.agentEngines.privateDelete(agentEngineName, true, null);
System.out.println("Delete operation started: " + deleteOperation.name().orElse(""));
}
runSandboxTemplatesExample(client);
runSandboxSnapshotsExample(client);
}

/** Demonstrates the sandbox templates submodule: create, get, list, and delete. */
private static void runSandboxTemplatesExample(Client client, String agentEngineName) {
private static void runSandboxTemplatesExample(Client client) {
System.out.println("\n=== Sandbox Templates ===");

// 1. Create a sandbox environment template.
// 1. Create a computer-use sandbox environment template (no agent engine name needed).
System.out.println("\n--- Creating a sandbox template ---");
SandboxEnvironmentTemplateOperation templateOp =
client.agentEngines.sandboxes.templates.privateCreate(
agentEngineName,
"Example Sandbox Template",
CreateSandboxEnvironmentTemplateConfig.builder()
.defaultContainerEnvironment(
SandboxEnvironmentTemplateDefaultContainerEnvironment.builder()
.defaultContainerCategory("DEFAULT_CONTAINER_CATEGORY_COMPUTER_USE")
.build())
.egressControlConfig(
SandboxEnvironmentTemplateEgressControlConfig.builder()
.internetAccess(true)
.build())
.build());
client.sandboxes.templates.create("Example Sandbox Template", computerUseTemplateConfig());
System.out.println("Create template operation: " + templateOp.name().orElse(""));

// Wait for the template creation to complete.
while (!templateOp.done().filter(Boolean::booleanValue).isPresent()) {
sleep(10000);
templateOp =
client.agentEngines.sandboxes.templates.getSandboxEnvironmentTemplateOperation(
client.sandboxes.templates.getSandboxEnvironmentTemplateOperation(
templateOp.name().get(), null);
}
String templateName = templateOp.response().get().name().get();
System.out.println("Created Sandbox Template: " + templateName);

// 2. Get the template.
System.out.println("\n--- Getting the sandbox template ---");
SandboxEnvironmentTemplate template =
client.agentEngines.sandboxes.templates.privateGet(templateName, null);
SandboxEnvironmentTemplate template = client.sandboxes.templates.get(templateName, null);
System.out.println("Template display name: " + template.displayName().orElse(""));

// 3. List templates.
// 3. List templates (no agent engine name needed).
System.out.println("\n--- Listing sandbox templates ---");
ListSandboxEnvironmentTemplatesResponse listResponse =
client.agentEngines.sandboxes.templates.privateList(agentEngineName, null);
ListSandboxEnvironmentTemplatesResponse listResponse = client.sandboxes.templates.list(null);
System.out.println(
"Templates found: " + listResponse.sandboxEnvironmentTemplates().map(List::size).orElse(0));

// 4. Delete the template.
System.out.println("\n--- Deleting the sandbox template ---");
DeleteSandboxEnvironmentTemplateOperation deleteOp =
client.agentEngines.sandboxes.templates.privateDelete(templateName, null);
client.sandboxes.templates.delete(templateName, null);
System.out.println("Delete template operation: " + deleteOp.name().orElse(""));
}

/**
* Demonstrates the sandbox snapshots submodule: create a sandbox, snapshot it, get, list, and
* delete.
* Demonstrates the sandbox environments and snapshots submodules: create a template-based
* sandbox, snapshot it, get, list, restore a new sandbox from the snapshot, then clean up.
*/
private static void runSandboxSnapshotsExample(Client client, String agentEngineName) {
private static void runSandboxSnapshotsExample(Client client) {
System.out.println("\n=== Sandbox Snapshots ===");

// A snapshot captures the state of an existing sandbox environment, so first create one.
System.out.println("\n--- Creating a source sandbox ---");
AgentEngineSandboxOperation sandboxOp =
client.agentEngines.sandboxes.privateCreate(
agentEngineName,
SandboxEnvironmentSpec.builder()
.codeExecutionEnvironment(
SandboxEnvironmentSpecCodeExecutionEnvironment.builder()
.machineConfig("MACHINE_CONFIG_VCPU4_RAM4GIB")
.build())
.build(),
null);
while (!sandboxOp.done().filter(Boolean::booleanValue).isPresent()) {
// Snapshots can only be taken of template-based sandboxes, so create a computer-use template
// and then a sandbox environment from it.
System.out.println("\n--- Creating a computer-use template ---");
SandboxEnvironmentTemplateOperation templateOp =
client.sandboxes.templates.create("Snapshot Source Template", computerUseTemplateConfig());
while (!templateOp.done().filter(Boolean::booleanValue).isPresent()) {
sleep(10000);
sandboxOp =
client.agentEngines.sandboxes.privateGetSandboxOperation(sandboxOp.name().get(), null);
templateOp =
client.sandboxes.templates.getSandboxEnvironmentTemplateOperation(
templateOp.name().get(), null);
}
String sandboxName = sandboxOp.response().get().name().get();
System.out.println("Created source Sandbox: " + sandboxName);
String templateName = templateOp.response().get().name().get();
System.out.println("Created Sandbox Template: " + templateName);

String sandboxName = null;
String restoredSandboxName = null;
try {
// 1. Create a snapshot of the sandbox.
// Create a sandbox environment from the template (no agent engine name needed).
System.out.println("\n--- Creating a sandbox environment from the template ---");
AgentEngineSandboxOperation sandboxOp =
client.sandboxes.environments.create(
/* spec= */ null,
CreateAgentEngineSandboxConfig.builder()
.displayName("Snapshot Source Sandbox")
.sandboxEnvironmentTemplate(templateName)
.build());
while (!sandboxOp.done().filter(Boolean::booleanValue).isPresent()) {
sleep(10000);
sandboxOp = client.sandboxes.environments.getSandboxOperation(sandboxOp.name().get(), null);
}
sandboxName = sandboxOp.response().get().name().get();
System.out.println("Created source Sandbox: " + sandboxName);

// 1. Create a snapshot of the running sandbox.
System.out.println("\n--- Creating a sandbox snapshot ---");
AgentEngineSandboxSnapshotOperation snapshotOp =
client.agentEngines.sandboxes.snapshots.privateCreate(
client.sandboxes.snapshots.create(
sandboxName,
CreateAgentEngineSandboxSnapshotConfig.builder()
.displayName("Example Sandbox Snapshot")
.ttl(Duration.ofSeconds(3600))
.build());
System.out.println("Create snapshot operation: " + snapshotOp.name().orElse(""));

// Wait for the snapshot creation to complete.
while (!snapshotOp.done().filter(Boolean::booleanValue).isPresent()) {
sleep(10000);
snapshotOp =
client.agentEngines.sandboxes.snapshots.getSandboxSnapshotOperation(
snapshotOp.name().get(), null);
client.sandboxes.snapshots.getSandboxSnapshotOperation(snapshotOp.name().get(), null);
}
String snapshotName = snapshotOp.response().get().name().get();
System.out.println("Created Sandbox Snapshot: " + snapshotName);

// 2. Get the snapshot.
System.out.println("\n--- Getting the sandbox snapshot ---");
SandboxEnvironmentSnapshot snapshot =
client.agentEngines.sandboxes.snapshots.privateGet(snapshotName, null);
SandboxEnvironmentSnapshot snapshot = client.sandboxes.snapshots.get(snapshotName, null);
System.out.println("Snapshot Name: " + snapshot.name().orElse(""));

// 3. List snapshots.
// 3. List snapshots (no agent engine name needed).
System.out.println("\n--- Listing sandbox snapshots ---");
ListSandboxEnvironmentSnapshotsResponse listResponse =
client.agentEngines.sandboxes.snapshots.privateList(agentEngineName, null);
ListSandboxEnvironmentSnapshotsResponse listResponse = client.sandboxes.snapshots.list(null);
System.out.println(
"Snapshots found: "
+ listResponse.sandboxEnvironmentSnapshots().map(List::size).orElse(0));

// 4. Delete the snapshot.
// 4. Restore a new sandbox environment from the snapshot.
System.out.println("\n--- Restoring a sandbox from the snapshot ---");
AgentEngineSandboxOperation restoreOp =
client.sandboxes.environments.create(
/* spec= */ null,
CreateAgentEngineSandboxConfig.builder()
.displayName("Restored Sandbox")
.sandboxEnvironmentSnapshot(snapshotName)
.build());
while (!restoreOp.done().filter(Boolean::booleanValue).isPresent()) {
sleep(10000);
restoreOp = client.sandboxes.environments.getSandboxOperation(restoreOp.name().get(), null);
}
restoredSandboxName = restoreOp.response().get().name().get();
System.out.println("Restored Sandbox: " + restoredSandboxName);

// 5. Delete the snapshot.
System.out.println("\n--- Deleting the sandbox snapshot ---");
DeleteSandboxEnvironmentSnapshotOperation deleteOp =
client.agentEngines.sandboxes.snapshots.privateDelete(snapshotName, null);
client.sandboxes.snapshots.delete(snapshotName, null);
System.out.println("Delete snapshot operation: " + deleteOp.name().orElse(""));
} finally {
// Clean up the source sandbox.
System.out.println("\n--- Deleting the source sandbox ---");
DeleteAgentEngineSandboxOperation deleteSandboxOp =
client.agentEngines.sandboxes.privateDelete(sandboxName, null);
System.out.println("Delete sandbox operation: " + deleteSandboxOp.name().orElse(""));
// Clean up the restored sandbox, source sandbox, and template.
if (restoredSandboxName != null) {
System.out.println("\n--- Deleting the restored sandbox ---");
client.sandboxes.environments.delete(restoredSandboxName, null);
}
if (sandboxName != null) {
System.out.println("\n--- Deleting the source sandbox ---");
DeleteAgentEngineSandboxOperation deleteSandboxOp =
client.sandboxes.environments.delete(sandboxName, null);
System.out.println("Delete sandbox operation: " + deleteSandboxOp.name().orElse(""));
}
System.out.println("\n--- Deleting the source template ---");
client.sandboxes.templates.delete(templateName, null);
}
}

/** Returns the config for a computer-use sandbox template with internet access. */
private static CreateSandboxEnvironmentTemplateConfig computerUseTemplateConfig() {
return CreateSandboxEnvironmentTemplateConfig.builder()
.defaultContainerEnvironment(
SandboxEnvironmentTemplateDefaultContainerEnvironment.builder()
.defaultContainerCategory("DEFAULT_CONTAINER_CATEGORY_COMPUTER_USE")
.build())
.egressControlConfig(
SandboxEnvironmentTemplateEgressControlConfig.builder().internetAccess(true).build())
.build();
}

private static void sleep(long millis) {
try {
Thread.sleep(millis);
Expand Down
49 changes: 49 additions & 0 deletions src/main/java/com/google/cloud/agentplatform/AgentEngines.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1401,4 +1401,53 @@ public AgentEngineOperation privateUpdate(String name, UpdateAgentEngineConfig c
return processResponseForPrivateUpdate(response, config);
}
}

private static final java.util.concurrent.ConcurrentHashMap<String, String>
DEFAULT_AGENT_ENGINE_BY_LOCATION = new java.util.concurrent.ConcurrentHashMap<>();

private static final String DEFAULT_AGENT_ENGINE_DISPLAY_NAME =
"default-agent-platform-sandbox-host";

/**
* Returns the resource name of a shared "default" agent engine for the client's project and
* location, creating it lazily if it does not exist. Backs sandbox calls where the caller does
* not supply an agent engine name, so repeated calls reuse one agent engine instead of creating
* many.
*/
static String ensureDefaultAgentEngine(ApiClient apiClient) {
String key = apiClient.project() + "/" + apiClient.location();
return DEFAULT_AGENT_ENGINE_BY_LOCATION.computeIfAbsent(
key, unused -> new AgentEngines(apiClient).resolveDefaultAgentEngine());
}

private String resolveDefaultAgentEngine() {
ListReasoningEnginesResponse listResponse = privateList(null);
if (listResponse.reasoningEngines().isPresent()) {
for (ReasoningEngine reasoningEngine : listResponse.reasoningEngines().get()) {
if (DEFAULT_AGENT_ENGINE_DISPLAY_NAME.equals(reasoningEngine.displayName().orElse(null))
&& reasoningEngine.name().isPresent()) {
return reasoningEngine.name().get();
}
}
}
AgentEngineOperation operation =
privateCreate(
CreateAgentEngineConfig.builder()
.displayName(DEFAULT_AGENT_ENGINE_DISPLAY_NAME)
.build());
while (!operation.done().filter(Boolean::booleanValue).isPresent()) {
try {
Thread.sleep(5000);
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
throw new IllegalStateException("Interrupted while creating default agent engine.", e);
}
operation = privateGetAgentOperation(operation.name().get(), null);
}
if (operation.error().isPresent()) {
throw new IllegalStateException(
"Failed to create default agent engine: " + operation.error().get());
}
return operation.response().get().name().get();
}
}
Loading
Loading