feat(genai): Sample/batch prediction

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

@@ -0,0 +1,82 @@
+// 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}`);
+  }
+
+  // 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,6 +13,9 @@
     "test": "c8 mocha -p -j 2 --timeout 2400000 test/*.test.js test/**/*.test.js"
   },
   "dependencies": {
+    "@google-cloud/bigquery": "^8.1.1",
+    "@google-cloud/storage": "^7.17.1",
+
     "@google/genai": "1.20.0",
     "axios": "^1.6.2",
     "google-auth-library": "^10.3.0",

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

@@ -0,0 +1,59 @@
+// 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 projectId = process.env.CAIP_PROJECT_ID;
+const location = 'us-central1';
+const {delay} = require('./util');
+const proxyquire = require('proxyquire');
+const {GoogleGenAI_Mock} = require('./batchprediction-utils');
+
+const sample = proxyquire(
+  '../batch-prediction/batchpredict-embeddings-with-gcs',
+  {
+    '@google/genai': {
+      GoogleGenAI: GoogleGenAI_Mock,
+    },
+  }
+);
+
+async function getGcsOutputUri() {
+  return {
+    uri: 'gs://mock/output',
+    async cleanup() {},
+  };
+}
+
+describe('batchpredict-with-gcs', () => {
+  it('should return the batch job state', async function () {
+    this.timeout(500000);
+    this.retries(4);
+    await delay(this.test);
+    const gcsOutput = await getGcsOutputUri();
+    try {
+      const output = await sample.runBatchPredictionJob(
+        gcsOutput.uri,
+        projectId,
+        location
+      );
+      assert.notEqual(output, undefined);
+    } finally {
+      await gcsOutput.cleanup();
+    }
+  });
+});

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

@@ -0,0 +1,57 @@
+// 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 projectId = process.env.CAIP_PROJECT_ID;
+const location = 'us-central1';
+const {delay} = require('./util');
+
+const proxyquire = require('proxyquire');
+const {GoogleGenAI_Mock} = require('./batchprediction-utils');
+
+const sample = proxyquire('../batch-prediction/batchpredict-with-bq', {
+  '@google/genai': {
+    GoogleGenAI: GoogleGenAI_Mock,
+  },
+});
+
+async function getBqOutputUri() {
+  return {
+    uri: 'gs://mock/output',
+    async cleanup() {},
+  };
+}
+
+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();
+    }
+  });
+});