feat(genai): Sample/batch prediction

genai/batch-prediction/batchpredict-embeddings-with-gcs.js

@@ -0,0 +1,86 @@
+// 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
+//
+//     https://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.
+
+'use strict';
+
+// [START googlegenaisdk_batchpredict_embeddings_with_gcs]
+const {GoogleGenAI} = require('@google/genai');
+
+const GOOGLE_CLOUD_PROJECT = process.env.GOOGLE_CLOUD_PROJECT;
+const GOOGLE_CLOUD_LOCATION =
+  process.env.GOOGLE_CLOUD_LOCATION || 'us-central1';
+const OUTPUT_URI = 'gs://your-bucket/your-prefix';
+
+async function runBatchPredictionJob(
+  outputUri = OUTPUT_URI,
+  projectId = GOOGLE_CLOUD_PROJECT,
+  location = GOOGLE_CLOUD_LOCATION
+) {
+  const client = new GoogleGenAI({
+    vertexai: true,
+    project: projectId,
+    location: location,
+    httpOptions: {
+      apiVersion: 'v1',
+    },
+  });
+
+  // See the documentation: https://googleapis.github.io/js-genai/release_docs/classes/batches.Batches.html
+  let job = await client.batches.create({
+    model: 'text-embedding-005',
+    // Source link: https://storage.cloud.google.com/cloud-samples-data/batch/prompt_for_batch_gemini_predict.jsonl
+    src: 'gs://cloud-samples-data/generative-ai/embeddings/embeddings_input.jsonl',
+    config: {
+      dest: outputUri,
+    },
+  });
+
+  console.log(`Job name: ${job.name}`);
+  console.log(`Job state: ${job.state}`);
+
+  // Example response:
+  //  Job name: projects/%PROJECT_ID%/locations/us-central1/batchPredictionJobs/9876453210000000000
+  //  Job state: JOB_STATE_PENDING
+
+  const completedStates = new Set([
+    'JOB_STATE_SUCCEEDED',
+    'JOB_STATE_FAILED',
+    'JOB_STATE_CANCELLED',
+    'JOB_STATE_PAUSED',
+  ]);
+
+  while (!completedStates.has(job.state)) {
+    await new Promise(resolve => setTimeout(resolve, 30000));
+    job = await client.batches.get({name: job.name});
+    console.log(`Job state: ${job.state}`);
+    if (job.state === 'JOB_STATE_FAILED') {
+      console.log(`Error: ${job.state}`);
+      break;
+    }
+  }
+
+  // Example response:
+  //  Job state: JOB_STATE_PENDING
+  //  Job state: JOB_STATE_RUNNING
+  //  Job state: JOB_STATE_RUNNING
+  //  ...
+  //  Job state: JOB_STATE_SUCCEEDED
+
+  return job.state;
+}
+// [END googlegenaisdk_batchpredict_embeddings_with_gcs]
+
+module.exports = {
+  runBatchPredictionJob,
+};

genai/batch-prediction/batchpredict-with-bq.js

@@ -0,0 +1,83 @@
+// 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
+//
+//     https://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.
+
+'use strict';
+
+// [START googlegenaisdk_batchpredict_with_bq]
+const {GoogleGenAI} = require('@google/genai');
+
+const GOOGLE_CLOUD_PROJECT = process.env.GOOGLE_CLOUD_PROJECT;
+const GOOGLE_CLOUD_LOCATION =
+  process.env.GOOGLE_CLOUD_LOCATION || 'us-central1';
+const OUTPUT_URI = 'bq://your-project.your_dataset.your_table';
+
+async function runBatchPredictionJob(
+  outputUri = OUTPUT_URI,
+  projectId = GOOGLE_CLOUD_PROJECT,
+  location = GOOGLE_CLOUD_LOCATION
+) {
+  const client = new GoogleGenAI({
+    vertexai: true,
+    project: projectId,
+    location: location,
+    httpOptions: {
+      apiVersion: 'v1',
+    },
+  });
+
+  // See the documentation: https://googleapis.github.io/js-genai/release_docs/classes/batches.Batches.html
+  let job = await client.batches.create({
+    // To use a tuned model, set the model param to your tuned model using the following format:
+    // model="projects/{PROJECT_ID}/locations/{LOCATION}/models/{MODEL_ID}"
+    model: 'gemini-2.5-flash',
+    src: 'bq://storage-samples.generative_ai.batch_requests_for_multimodal_input',
+    config: {
+      dest: outputUri,
+    },
+  });
+
+  console.log(`Job name: ${job.name}`);
+  console.log(`Job state: ${job.state}`);
+
+  // Example response:
+  //  Job name: projects/%PROJECT_ID%/locations/us-central1/batchPredictionJobs/9876453210000000000
+  //  Job state: JOB_STATE_PENDING
+
+  const completedStates = new Set([
+    'JOB_STATE_SUCCEEDED',
+    'JOB_STATE_FAILED',
+    'JOB_STATE_CANCELLED',
+    'JOB_STATE_PAUSED',
+  ]);
+
+  while (!completedStates.has(job.state)) {
+    await new Promise(resolve => setTimeout(resolve, 30000));
+    job = await client.batches.get({name: job.name});
+    console.log(`Job state: ${job.state}`);
+  }
+
+  // Example response:
+  //  Job state: JOB_STATE_PENDING
+  //  Job state: JOB_STATE_RUNNING
+  //  Job state: JOB_STATE_RUNNING
+  //  ...
+  //  Job state: JOB_STATE_SUCCEEDED
+
+  return job.state;
+}
+// [END googlegenaisdk_batchpredict_with_bq]
+
+module.exports = {
+  runBatchPredictionJob,
+};

genai/batch-prediction/batchpredict-with-gcs.js

@@ -0,0 +1,84 @@
+// 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
+//
+//     https://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.
+
+'use strict';
+
+// [START googlegenaisdk_batchpredict_with_gcs]
+const {GoogleGenAI} = require('@google/genai');
+
+const GOOGLE_CLOUD_PROJECT = process.env.GOOGLE_CLOUD_PROJECT;
+const GOOGLE_CLOUD_LOCATION =
+  process.env.GOOGLE_CLOUD_LOCATION || 'us-central1';
+const OUTPUT_URI = 'gs://your-bucket/your-prefix';
+
+async function runBatchPredictionJob(
+  outputUri = OUTPUT_URI,
+  projectId = GOOGLE_CLOUD_PROJECT,
+  location = GOOGLE_CLOUD_LOCATION
+) {
+  const client = new GoogleGenAI({
+    vertexai: true,
+    project: projectId,
+    location: location,
+    httpOptions: {
+      apiVersion: 'v1',
+    },
+  });
+
+  // See the documentation: https://googleapis.github.io/js-genai/release_docs/classes/batches.Batches.html
+  let job = await client.batches.create({
+    // To use a tuned model, set the model param to your tuned model using the following format:
+    // model="projects/{PROJECT_ID}/locations/{LOCATION}/models/{MODEL_ID}"
+    model: 'gemini-2.5-flash',
+    // Source link: https://storage.cloud.google.com/cloud-samples-data/batch/prompt_for_batch_gemini_predict.jsonl
+    src: 'gs://cloud-samples-data/batch/prompt_for_batch_gemini_predict.jsonl',
+    config: {
+      dest: outputUri,
+    },
+  });
+
+  console.log(`Job name: ${job.name}`);
+  console.log(`Job state: ${job.state}`);
+
+  // Example response:
+  //  Job name: projects/%PROJECT_ID%/locations/us-central1/batchPredictionJobs/9876453210000000000
+  //  Job state: JOB_STATE_PENDING
+
+  const completedStates = new Set([
+    'JOB_STATE_SUCCEEDED',
+    'JOB_STATE_FAILED',
+    'JOB_STATE_CANCELLED',
+    'JOB_STATE_PAUSED',
+  ]);
+
+  while (!completedStates.has(job.state)) {
+    await new Promise(resolve => setTimeout(resolve, 30000));
+    job = await client.batches.get({name: job.name});
+    console.log(`Job state: ${job.state}`);
+  }
+
+  // Example response:
+  //  Job state: JOB_STATE_PENDING
+  //  Job state: JOB_STATE_RUNNING
+  //  Job state: JOB_STATE_RUNNING
+  //  ...
+  //  Job state: JOB_STATE_SUCCEEDED
+
+  return job.state;
+}
+// [END googlegenaisdk_batchpredict_with_gcs]
+
+module.exports = {
+  runBatchPredictionJob,
+};

genai/package.json

@@ -13,7 +13,9 @@
     "test": "c8 mocha -p -j 2 --timeout 2400000 test/*.test.js test/**/*.test.js"
   },
   "dependencies": {
-    "@google/genai": "1.12.0",
+    "@google-cloud/bigquery": "^8.1.1",
+    "@google-cloud/storage": "^7.17.1",
+    "@google/genai": "1.20.0",
     "axios": "^1.6.2",
     "luxon": "^3.7.1",
     "supertest": "^7.0.0"

genai/test/batchpredict-embeddings-with-gcs.test.js

@@ -0,0 +1,65 @@
+// 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
+//
+//     https://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.
+
+'use strict';
+
+const {assert} = require('chai');
+const {describe, it} = require('mocha');
+const {Storage} = require('@google-cloud/storage');
+
+const storage = new Storage();
+
+const GCS_OUTPUT_BUCKET = 'nodejs-docs-samples-tests';
+
+const projectId = process.env.CAIP_PROJECT_ID;
+const location = 'us-central1';
+const sample = require('../batch-prediction/batchpredict-embeddings-with-gcs');
+const {delay} = require('./util');
+
+async function getGcsOutputUri() {
+  const dt = new Date();
+  const prefix = `text_output/${dt.toISOString()}`;
+  const fullUri = `gs://${GCS_OUTPUT_BUCKET}/${prefix}`;
+
+  return {
+    uri: fullUri,
+    async cleanup() {
+      const [files] = await storage.bucket(GCS_OUTPUT_BUCKET).getFiles({
+        prefix,
+      });
+      for (const file of files) {
+        await file.delete();
+      }
+    },
+  };
+}
+
+describe('batchpredict-with-gcs', () => {
+  it('should return the batch job state', async function () {
+    this.timeout(500000);
+    this.retries(4);
+    await delay(this.test);
+    const bqOutput = await getGcsOutputUri();
+    try {
+      const output = await sample.runBatchPredictionJob(
+        bqOutput.uri,
+        projectId,
+        location
+      );
+      assert.notEqual(output, undefined);
+    } finally {
+      await bqOutput.cleanup();
+    }
+  });
+});

genai/test/batchpredict-with-bq.test.js

@@ -0,0 +1,64 @@
+// 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
+//
+//     https://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.
+
+'use strict';
+
+const {assert} = require('chai');
+const {describe, it} = require('mocha');
+const {BigQuery} = require('@google-cloud/bigquery');
+
+const bigquery = new BigQuery();
+
+const BQ_OUTPUT_DATASET = `${process.env.BQ_OUTPUT_DATASET}.gen_ai_batch_prediction`;
+
+const projectId = process.env.CAIP_PROJECT_ID;
+const location = 'us-central1';
+const {delay} = require('./util');
+
+const sample = require('../batch-prediction/batchpredict-with-bq');
+
+async function getBqOutputUri() {
+  const dt = new Date();
+  const tableName = `text_output_${dt.getFullYear()}_${dt.getMonth() + 1}_${dt.getDate()}_T${dt.getHours()}_${dt.getMinutes()}_${dt.getSeconds()}`;
+  const tableUri = `${BQ_OUTPUT_DATASET}.${tableName}`;
+  const fullUri = `bq://${tableUri}`;
+
+  return {
+    uri: fullUri,
+    async cleanup() {
+      await bigquery.dataset(BQ_OUTPUT_DATASET).table(tableName).delete({
+        ignoreNotFound: true,
+      });
+    },
+  };
+}
+
+describe('batchpredict-with-bq', () => {
+  it('should return the batch job state', async function () {
+    this.timeout(500000);
+    this.retries(4);
+    await delay(this.test);
+    const bqOutput = await getBqOutputUri();
+    try {
+      const output = await sample.runBatchPredictionJob(
+        bqOutput.uri,
+        projectId,
+        location
+      );
+      assert.notEqual(output, undefined);
+    } finally {
+      await bqOutput.cleanup();
+    }
+  });
+});