diff --git a/genai/snippets/pom.xml b/genai/snippets/pom.xml index 32015ec0bf4..585d72061cc 100644 --- a/genai/snippets/pom.xml +++ b/genai/snippets/pom.xml @@ -51,7 +51,11 @@ com.google.genai google-genai - 1.15.0 + 1.20.0 + + + com.google.cloud + google-cloud-storage junit diff --git a/genai/snippets/src/main/java/genai/batchprediction/BatchPredictionEmbeddingsWithGcs.java b/genai/snippets/src/main/java/genai/batchprediction/BatchPredictionEmbeddingsWithGcs.java new file mode 100644 index 00000000000..22519cc0bb1 --- /dev/null +++ b/genai/snippets/src/main/java/genai/batchprediction/BatchPredictionEmbeddingsWithGcs.java @@ -0,0 +1,105 @@ +/* + * 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.GetBatchJobConfig; +import com.google.genai.types.HttpOptions; +import com.google.genai.types.JobState; +import java.util.EnumSet; +import java.util.Optional; +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 Optional 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("Missing job name")); + System.out.println("Job name: " + jobName); + Optional jobState = batchJob.state(); + jobState.ifPresent(state -> System.out.println("Job state: " + state)); + // 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 completedStates = + EnumSet.of(JOB_STATE_SUCCEEDED, JOB_STATE_FAILED, JOB_STATE_CANCELLED, JOB_STATE_PAUSED); + + while (jobState.isPresent() && !completedStates.contains(jobState.get().knownEnum())) { + TimeUnit.SECONDS.sleep(30); + batchJob = client.batches.get(jobName, GetBatchJobConfig.builder().build()); + jobState = batchJob.state(); + batchJob.state().ifPresent(state -> System.out.println("Job state: " + state)); + } + // Example response: + // Job state: JOB_STATE_QUEUED + // Job state: JOB_STATE_RUNNING + // ... + // Job state: JOB_STATE_SUCCEEDED + return jobState; + } + } +} +// [END googlegenaisdk_batchpredict_embeddings_with_gcs] diff --git a/genai/snippets/src/main/java/genai/batchprediction/BatchPredictionWithGcs.java b/genai/snippets/src/main/java/genai/batchprediction/BatchPredictionWithGcs.java new file mode 100644 index 00000000000..675038a6de7 --- /dev/null +++ b/genai/snippets/src/main/java/genai/batchprediction/BatchPredictionWithGcs.java @@ -0,0 +1,107 @@ +/* + * 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.GetBatchJobConfig; +import com.google.genai.types.HttpOptions; +import com.google.genai.types.JobState; +import java.util.EnumSet; +import java.util.Optional; +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 Optional 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("Missing job name")); + System.out.println("Job name: " + jobName); + Optional jobState = batchJob.state(); + jobState.ifPresent(state -> System.out.println("Job state: " + state)); + // 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 completedStates = + EnumSet.of(JOB_STATE_SUCCEEDED, JOB_STATE_FAILED, JOB_STATE_CANCELLED, JOB_STATE_PAUSED); + + while (jobState.isPresent() && !completedStates.contains(jobState.get().knownEnum())) { + TimeUnit.SECONDS.sleep(30); + batchJob = client.batches.get(jobName, GetBatchJobConfig.builder().build()); + jobState = batchJob.state(); + batchJob.state().ifPresent(state -> System.out.println("Job state: " + state)); + } + // 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] diff --git a/genai/snippets/src/test/java/genai/batchprediction/BatchPredictionIT.java b/genai/snippets/src/test/java/genai/batchprediction/BatchPredictionIT.java new file mode 100644 index 00000000000..b0e2c46f6b0 --- /dev/null +++ b/genai/snippets/src/test/java/genai/batchprediction/BatchPredictionIT.java @@ -0,0 +1,106 @@ +/* + * 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.Optional; +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 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 { + Optional response = + BatchPredictionWithGcs.createBatchJob(GEMINI_FLASH, OUTPUT_GCS_URI); + assertThat(response).isPresent(); + assertThat(response.get().toString()).isNotEmpty(); + assertThat(response.get().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 { + Optional response = + BatchPredictionEmbeddingsWithGcs.createBatchJob(EMBEDDING_MODEL, OUTPUT_GCS_URI); + assertThat(response).isPresent(); + assertThat(response.get().toString()).isNotEmpty(); + assertThat(response.get().toString()).isEqualTo("JOB_STATE_SUCCEEDED"); + assertThat(bout.toString()).contains("Job name: "); + assertThat(bout.toString()).contains("Job state: JOB_STATE_SUCCEEDED"); + } +}