feat(genai): add batch prediction samples (1) and update GenAI SDK

genai/snippets/pom.xml

@@ -51,7 +51,11 @@
         <dependency>
             <groupId>com.google.genai</groupId>
             <artifactId>google-genai</artifactId>
-            <version>1.15.0</version>
+            <version>1.20.0</version>
+        </dependency>
+        <dependency>
+            <groupId>com.google.cloud</groupId>
+            <artifactId>google-cloud-storage</artifactId>
         </dependency>
         <dependency>
             <artifactId>junit</artifactId>

genai/snippets/src/main/java/genai/batchprediction/BatchPredictionEmbeddingsWithGcs.java

@@ -0,0 +1,110 @@
+/*
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package genai.batchprediction;
+
+// [START googlegenaisdk_batchpredict_embeddings_with_gcs]
+
+import static com.google.genai.types.JobState.Known.JOB_STATE_CANCELLED;
+import static com.google.genai.types.JobState.Known.JOB_STATE_FAILED;
+import static com.google.genai.types.JobState.Known.JOB_STATE_PAUSED;
+import static com.google.genai.types.JobState.Known.JOB_STATE_SUCCEEDED;
+
+import com.google.genai.Client;
+import com.google.genai.types.BatchJob;
+import com.google.genai.types.BatchJobDestination;
+import com.google.genai.types.BatchJobSource;
+import com.google.genai.types.CreateBatchJobConfig;
+import com.google.genai.types.HttpOptions;
+import com.google.genai.types.JobState;
+import java.util.EnumSet;
+import java.util.Set;
+import java.util.concurrent.TimeUnit;
+
+public class BatchPredictionEmbeddingsWithGcs {
+
+  public static void main(String[] args) throws InterruptedException {
+    // TODO(developer): Replace these variables before running the sample.
+    String modelId = "text-embedding-005";
+    String outputGcsUri = "gs://your-bucket/your-prefix";
+    createBatchJob(modelId, outputGcsUri);
+  }
+
+  // Creates a batch prediction job with embedding model and Google Cloud Storage
+  public static JobState createBatchJob(String modelId, String outputGcsUri)
+      throws InterruptedException {
+    // Client Initialization. Once created, it can be reused for multiple requests.
+    try (Client client =
+        Client.builder()
+            .location("us-central1")
+            .vertexAI(true)
+            .httpOptions(HttpOptions.builder().apiVersion("v1").build())
+            .build()) {
+
+      // See the documentation:
+      // https://googleapis.github.io/java-genai/javadoc/com/google/genai/Batches.html
+      BatchJobSource batchJobSource =
+          BatchJobSource.builder()
+              // Source link:
+              // https://storage.cloud.google.com/cloud-samples-data/generative-ai/embeddings/embeddings_input.jsonl
+              .gcsUri("gs://cloud-samples-data/generative-ai/embeddings/embeddings_input.jsonl")
+              .format("jsonl")
+              .build();
+
+      CreateBatchJobConfig batchJobConfig =
+          CreateBatchJobConfig.builder()
+              .displayName("your-display-name")
+              .dest(BatchJobDestination.builder().gcsUri(outputGcsUri).format("jsonl").build())
+              .build();
+
+      BatchJob batchJob = client.batches.create(modelId, batchJobSource, batchJobConfig);
+
+      String jobName =
+          batchJob.name().orElseThrow(() -> new IllegalStateException("Failed to get job name."));
+      JobState jobState =
+          batchJob.state().orElseThrow(() -> new IllegalStateException("Failed to get job state."));
+
+      System.out.println("Job name: " + jobName);
+      System.out.println("Job state: " + jobState);
+      // Job name:
+      // projects/{PROJECT_ID}/locations/us-central1/batchPredictionJobs/6205497615459549184
+      // Job state: JOB_STATE_PENDING
+
+      // See the documentation:
+      // https://googleapis.github.io/java-genai/javadoc/com/google/genai/types/BatchJob.html
+      Set<JobState.Known> completedStates =
+          EnumSet.of(JOB_STATE_SUCCEEDED, JOB_STATE_FAILED, JOB_STATE_CANCELLED, JOB_STATE_PAUSED);
+
+      while (!completedStates.contains(jobState.knownEnum())) {
+        TimeUnit.SECONDS.sleep(30);
+        batchJob = client.batches.get(jobName, null);
+        jobState =
+            batchJob
+                .state()
+                .orElseThrow(() -> new IllegalStateException("Failed to get job state."));
+        System.out.println("Job state: " + jobState);
+      }
+      // Example response:
+      // Job state: JOB_STATE_QUEUED
+      // Job state: JOB_STATE_RUNNING
+      // Job state: JOB_STATE_RUNNING
+      // ...
+      // Job state: JOB_STATE_SUCCEEDED
+      return jobState;
+    }
+  }
+}
+// [END googlegenaisdk_batchpredict_embeddings_with_gcs]

genai/snippets/src/main/java/genai/batchprediction/BatchPredictionWithGcs.java

@@ -0,0 +1,112 @@
+/*
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package genai.batchprediction;
+
+// [START googlegenaisdk_batchpredict_with_gcs]
+
+import static com.google.genai.types.JobState.Known.JOB_STATE_CANCELLED;
+import static com.google.genai.types.JobState.Known.JOB_STATE_FAILED;
+import static com.google.genai.types.JobState.Known.JOB_STATE_PAUSED;
+import static com.google.genai.types.JobState.Known.JOB_STATE_SUCCEEDED;
+
+import com.google.genai.Client;
+import com.google.genai.types.BatchJob;
+import com.google.genai.types.BatchJobDestination;
+import com.google.genai.types.BatchJobSource;
+import com.google.genai.types.CreateBatchJobConfig;
+import com.google.genai.types.HttpOptions;
+import com.google.genai.types.JobState;
+import java.util.EnumSet;
+import java.util.Set;
+import java.util.concurrent.TimeUnit;
+
+public class BatchPredictionWithGcs {
+
+  public static void main(String[] args) throws InterruptedException {
+    // TODO(developer): Replace these variables before running the sample.
+    // To use a tuned model, set the model param to your tuned model using the following format:
+    // modelId = "projects/{PROJECT_ID}/locations/{LOCATION}/models/{MODEL_ID}
+    String modelId = "gemini-2.5-flash";
+    String outputGcsUri = "gs://your-bucket/your-prefix";
+    createBatchJob(modelId, outputGcsUri);
+  }
+
+  // Creates a batch prediction job with Google Cloud Storage
+  public static JobState createBatchJob(String modelId, String outputGcsUri)
+      throws InterruptedException {
+    // Client Initialization. Once created, it can be reused for multiple requests.
+    try (Client client =
+        Client.builder()
+            .location("us-central1")
+            .vertexAI(true)
+            .httpOptions(HttpOptions.builder().apiVersion("v1").build())
+            .build()) {
+      // See the documentation:
+      // https://googleapis.github.io/java-genai/javadoc/com/google/genai/Batches.html
+      BatchJobSource batchJobSource =
+          BatchJobSource.builder()
+              // Source link:
+              // https://storage.cloud.google.com/cloud-samples-data/batch/prompt_for_batch_gemini_predict.jsonl
+              .gcsUri("gs://cloud-samples-data/batch/prompt_for_batch_gemini_predict.jsonl")
+              .format("jsonl")
+              .build();
+
+      CreateBatchJobConfig batchJobConfig =
+          CreateBatchJobConfig.builder()
+              .displayName("your-display-name")
+              .dest(BatchJobDestination.builder().gcsUri(outputGcsUri).format("jsonl").build())
+              .build();
+
+      BatchJob batchJob = client.batches.create(modelId, batchJobSource, batchJobConfig);
+
+      String jobName =
+          batchJob.name().orElseThrow(() -> new IllegalStateException("Failed to get job name."));
+      JobState jobState =
+          batchJob.state().orElseThrow(() -> new IllegalStateException("Failed to get job state."));
+
+      System.out.println("Job name: " + jobName);
+      System.out.println("Job state: " + jobState);
+      // Example response:
+      // Job name:
+      // projects/{PROJECT_ID}/locations/us-central1/batchPredictionJobs/6205497615459549184
+      // Job state: JOB_STATE_PENDING
+
+      // See the documentation:
+      // https://googleapis.github.io/java-genai/javadoc/com/google/genai/types/BatchJob.html
+      Set<JobState.Known> completedStates =
+          EnumSet.of(JOB_STATE_SUCCEEDED, JOB_STATE_FAILED, JOB_STATE_CANCELLED, JOB_STATE_PAUSED);
+
+      while (!completedStates.contains(jobState.knownEnum())) {
+        TimeUnit.SECONDS.sleep(30);
+        batchJob = client.batches.get(jobName, null);
+        jobState =
+            batchJob
+                .state()
+                .orElseThrow(() -> new IllegalStateException("Failed to get job state."));
+        System.out.println("Job state: " + jobState);
+      }
+      // Example response:
+      // Job state: JOB_STATE_QUEUED
+      // Job state: JOB_STATE_RUNNING
+      // Job state: JOB_STATE_RUNNING
+      // ...
+      // Job state: JOB_STATE_SUCCEEDED
+      return jobState;
+    }
+  }
+}
+// [END googlegenaisdk_batchpredict_with_gcs]

genai/snippets/src/test/java/genai/batchprediction/BatchPredictionIT.java

@@ -0,0 +1,102 @@
+/*
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package genai.batchprediction;
+
+import static com.google.common.truth.Truth.assertThat;
+import static com.google.common.truth.Truth.assertWithMessage;
+
+import com.google.api.gax.paging.Page;
+import com.google.cloud.storage.Blob;
+import com.google.cloud.storage.Storage;
+import com.google.cloud.storage.StorageOptions;
+import com.google.genai.types.JobState;
+import java.io.ByteArrayOutputStream;
+import java.io.PrintStream;
+import java.util.UUID;
+import org.junit.After;
+import org.junit.AfterClass;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class BatchPredictionIT {
+
+  private static final String GEMINI_FLASH = "gemini-2.5-flash";
+  private static final String EMBEDDING_MODEL = "text-embedding-005";
+  private static final String BUCKET_NAME = "java-docs-samples-testing";
+  private static final String PREFIX = "genai-batch-prediction" + UUID.randomUUID();
+  private static final String OUTPUT_GCS_URI = String.format("gs://%s/%s", BUCKET_NAME, PREFIX);
+  private ByteArrayOutputStream bout;
+  private PrintStream out;
+
+  // Check if the required environment variables are set.
+  public static void requireEnvVar(String envVarName) {
+    assertWithMessage(String.format("Missing environment variable '%s' ", envVarName))
+        .that(System.getenv(envVarName))
+        .isNotEmpty();
+  }
+
+  @BeforeClass
+  public static void checkRequirements() {
+    requireEnvVar("GOOGLE_CLOUD_PROJECT");
+  }
+
+  @AfterClass
+  public static void cleanup() {
+    Storage storage = StorageOptions.getDefaultInstance().getService();
+    Page<Blob> blobs = storage.list(BUCKET_NAME, Storage.BlobListOption.prefix(PREFIX));
+
+    for (Blob blob : blobs.iterateAll()) {
+      storage.delete(blob.getBlobId());
+    }
+  }
+
+  @Before
+  public void setUp() {
+    bout = new ByteArrayOutputStream();
+    out = new PrintStream(bout);
+    System.setOut(out);
+  }
+
+  @After
+  public void tearDown() {
+    System.setOut(null);
+    bout.reset();
+  }
+
+  @Test
+  public void testBatchPredictionWithGcs() throws InterruptedException {
+    JobState response = BatchPredictionWithGcs.createBatchJob(GEMINI_FLASH, OUTPUT_GCS_URI);
+    assertThat(response.toString()).isNotEmpty();
+    assertThat(response.toString()).isEqualTo("JOB_STATE_SUCCEEDED");
+    assertThat(bout.toString()).contains("Job name: ");
+    assertThat(bout.toString()).contains("Job state: JOB_STATE_SUCCEEDED");
+  }
+
+  @Test
+  public void testBatchPredictionEmbeddingsWithGcs() throws InterruptedException {
+    JobState response =
+        BatchPredictionEmbeddingsWithGcs.createBatchJob(EMBEDDING_MODEL, OUTPUT_GCS_URI);
+    assertThat(response.toString()).isNotEmpty();
+    assertThat(response.toString()).isEqualTo("JOB_STATE_SUCCEEDED");
+    assertThat(bout.toString()).contains("Job name: ");
+    assertThat(bout.toString()).contains("Job state: JOB_STATE_SUCCEEDED");
+  }
+}