Skip to content

Commit 77b05d8

Browse files
Guinersholtskinner
Guiners
and
holtskinner
authored
feat(genai): Sample/batch prediction (#4181)
Co-authored-by: Holt Skinner <13262395+holtskinner@users.noreply.github.com>
1 parent 9d1b80e commit 77b05d8

9 files changed

+498
-1
lines changed

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

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
// Copyright 2025 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// https://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
'use strict';
16+
17+
// [START googlegenaisdk_batchpredict_embeddings_with_gcs]
18+
const {GoogleGenAI} = require('@google/genai');
19+
20+
const GOOGLE_CLOUD_PROJECT = process.env.GOOGLE_CLOUD_PROJECT;
21+
const GOOGLE_CLOUD_LOCATION =
22+
process.env.GOOGLE_CLOUD_LOCATION || 'us-central1';
23+
const OUTPUT_URI = 'gs://your-bucket/your-prefix';
24+
25+
async function runBatchPredictionJob(
26+
outputUri = OUTPUT_URI,
27+
projectId = GOOGLE_CLOUD_PROJECT,
28+
location = GOOGLE_CLOUD_LOCATION
29+
) {
30+
const client = new GoogleGenAI({
31+
vertexai: true,
32+
project: projectId,
33+
location: location,
34+
httpOptions: {
35+
apiVersion: 'v1',
36+
},
37+
});
38+
39+
// See the documentation: https://googleapis.github.io/js-genai/release_docs/classes/batches.Batches.html
40+
let job = await client.batches.create({
41+
model: 'text-embedding-005',
42+
// Source link: https://storage.cloud.google.com/cloud-samples-data/batch/prompt_for_batch_gemini_predict.jsonl
43+
src: 'gs://cloud-samples-data/generative-ai/embeddings/embeddings_input.jsonl',
44+
config: {
45+
dest: outputUri,
46+
},
47+
});
48+
49+
console.log(`Job name: ${job.name}`);
50+
console.log(`Job state: ${job.state}`);
51+
52+
// Example response:
53+
// Job name: projects/%PROJECT_ID%/locations/us-central1/batchPredictionJobs/9876453210000000000
54+
// Job state: JOB_STATE_PENDING
55+
56+
const completedStates = new Set([
57+
'JOB_STATE_SUCCEEDED',
58+
'JOB_STATE_FAILED',
59+
'JOB_STATE_CANCELLED',
60+
'JOB_STATE_PAUSED',
61+
]);
62+
63+
while (!completedStates.has(job.state)) {
64+
await new Promise(resolve => setTimeout(resolve, 30000));
65+
job = await client.batches.get({name: job.name});
66+
console.log(`Job state: ${job.state}`);
67+
}
68+
69+
// Example response:
70+
// Job state: JOB_STATE_PENDING
71+
// Job state: JOB_STATE_RUNNING
72+
// Job state: JOB_STATE_RUNNING
73+
// ...
74+
// Job state: JOB_STATE_SUCCEEDED
75+
76+
return job.state;
77+
}
78+
// [END googlegenaisdk_batchpredict_embeddings_with_gcs]
79+
80+
module.exports = {
81+
runBatchPredictionJob,
82+
};

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

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
// Copyright 2025 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// https://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
'use strict';
16+
17+
// [START googlegenaisdk_batchpredict_with_bq]
18+
const {GoogleGenAI} = require('@google/genai');
19+
20+
const GOOGLE_CLOUD_PROJECT = process.env.GOOGLE_CLOUD_PROJECT;
21+
const GOOGLE_CLOUD_LOCATION =
22+
process.env.GOOGLE_CLOUD_LOCATION || 'us-central1';
23+
const OUTPUT_URI = 'bq://your-project.your_dataset.your_table';
24+
25+
async function runBatchPredictionJob(
26+
outputUri = OUTPUT_URI,
27+
projectId = GOOGLE_CLOUD_PROJECT,
28+
location = GOOGLE_CLOUD_LOCATION
29+
) {
30+
const client = new GoogleGenAI({
31+
vertexai: true,
32+
project: projectId,
33+
location: location,
34+
httpOptions: {
35+
apiVersion: 'v1',
36+
},
37+
});
38+
39+
// See the documentation: https://googleapis.github.io/js-genai/release_docs/classes/batches.Batches.html
40+
let job = await client.batches.create({
41+
// To use a tuned model, set the model param to your tuned model using the following format:
42+
// model="projects/{PROJECT_ID}/locations/{LOCATION}/models/{MODEL_ID}"
43+
model: 'gemini-2.5-flash',
44+
src: 'bq://storage-samples.generative_ai.batch_requests_for_multimodal_input',
45+
config: {
46+
dest: outputUri,
47+
},
48+
});
49+
50+
console.log(`Job name: ${job.name}`);
51+
console.log(`Job state: ${job.state}`);
52+
53+
// Example response:
54+
// Job name: projects/%PROJECT_ID%/locations/us-central1/batchPredictionJobs/9876453210000000000
55+
// Job state: JOB_STATE_PENDING
56+
57+
const completedStates = new Set([
58+
'JOB_STATE_SUCCEEDED',
59+
'JOB_STATE_FAILED',
60+
'JOB_STATE_CANCELLED',
61+
'JOB_STATE_PAUSED',
62+
]);
63+
64+
while (!completedStates.has(job.state)) {
65+
await new Promise(resolve => setTimeout(resolve, 30000));
66+
job = await client.batches.get({name: job.name});
67+
console.log(`Job state: ${job.state}`);
68+
}
69+
70+
// Example response:
71+
// Job state: JOB_STATE_PENDING
72+
// Job state: JOB_STATE_RUNNING
73+
// Job state: JOB_STATE_RUNNING
74+
// ...
75+
// Job state: JOB_STATE_SUCCEEDED
76+
77+
return job.state;
78+
}
79+
// [END googlegenaisdk_batchpredict_with_bq]
80+
81+
module.exports = {
82+
runBatchPredictionJob,
83+
};

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

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
// Copyright 2025 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// https://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
'use strict';
16+
17+
// [START googlegenaisdk_batchpredict_with_gcs]
18+
const {GoogleGenAI} = require('@google/genai');
19+
20+
const GOOGLE_CLOUD_PROJECT = process.env.GOOGLE_CLOUD_PROJECT;
21+
const GOOGLE_CLOUD_LOCATION =
22+
process.env.GOOGLE_CLOUD_LOCATION || 'us-central1';
23+
const OUTPUT_URI = 'gs://your-bucket/your-prefix';
24+
25+
async function runBatchPredictionJob(
26+
outputUri = OUTPUT_URI,
27+
projectId = GOOGLE_CLOUD_PROJECT,
28+
location = GOOGLE_CLOUD_LOCATION
29+
) {
30+
const client = new GoogleGenAI({
31+
vertexai: true,
32+
project: projectId,
33+
location: location,
34+
httpOptions: {
35+
apiVersion: 'v1',
36+
},
37+
});
38+
39+
// See the documentation: https://googleapis.github.io/js-genai/release_docs/classes/batches.Batches.html
40+
let job = await client.batches.create({
41+
// To use a tuned model, set the model param to your tuned model using the following format:
42+
// model="projects/{PROJECT_ID}/locations/{LOCATION}/models/{MODEL_ID}"
43+
model: 'gemini-2.5-flash',
44+
// Source link: https://storage.cloud.google.com/cloud-samples-data/batch/prompt_for_batch_gemini_predict.jsonl
45+
src: 'gs://cloud-samples-data/batch/prompt_for_batch_gemini_predict.jsonl',
46+
config: {
47+
dest: outputUri,
48+
},
49+
});
50+
51+
console.log(`Job name: ${job.name}`);
52+
console.log(`Job state: ${job.state}`);
53+
54+
// Example response:
55+
// Job name: projects/%PROJECT_ID%/locations/us-central1/batchPredictionJobs/9876453210000000000
56+
// Job state: JOB_STATE_PENDING
57+
58+
const completedStates = new Set([
59+
'JOB_STATE_SUCCEEDED',
60+
'JOB_STATE_FAILED',
61+
'JOB_STATE_CANCELLED',
62+
'JOB_STATE_PAUSED',
63+
]);
64+
65+
while (!completedStates.has(job.state)) {
66+
await new Promise(resolve => setTimeout(resolve, 30000));
67+
job = await client.batches.get({name: job.name});
68+
console.log(`Job state: ${job.state}`);
69+
}
70+
71+
// Example response:
72+
// Job state: JOB_STATE_PENDING
73+
// Job state: JOB_STATE_RUNNING
74+
// Job state: JOB_STATE_RUNNING
75+
// ...
76+
// Job state: JOB_STATE_SUCCEEDED
77+
78+
return job.state;
79+
}
80+
// [END googlegenaisdk_batchpredict_with_gcs]
81+
82+
module.exports = {
83+
runBatchPredictionJob,
84+
};

‎genai/batch-prediction/package.json‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
{
2+
"name": "nodejs-genai-batch-prediction-samples",
3+
"version": "0.0.1",
4+
"private": true,
5+
"license": "Apache-2.0",
6+
"author": "Google LLC",
7+
"repository": {
8+
"type": "git",
9+
"url": "https://github.com/GoogleCloudPlatform/nodejs-docs-samples.git"
10+
},
11+
"engines": {
12+
"node": ">=16.0.0"
13+
},
14+
"files": [
15+
"*.js"
16+
],
17+
"scripts": {
18+
"test": "c8 mocha -p -j 2 --timeout 2400000 test/*.test.js test/**/*.test.js"
19+
},
20+
"dependencies": {
21+
"@google/genai": "1.30.0"
22+
},
23+
"devDependencies": {
24+
"c8": "^10.0.0",
25+
"chai": "^4.5.0",
26+
"mocha": "^10.0.0",
27+
"proxyquire": "^2.1.3"
28+
}
29+
}

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

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
// Copyright 2025 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// https://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
'use strict';
16+
17+
const {assert} = require('chai');
18+
const {describe, it} = require('mocha');
19+
20+
const projectId = process.env.CAIP_PROJECT_ID;
21+
const location = 'us-central1';
22+
const {delay} = require('../../test/util');
23+
const proxyquire = require('proxyquire');
24+
const {GoogleGenAI_Mock} = require('./batchprediction-utils');
25+
26+
const sample = proxyquire('../batchpredict-embeddings-with-gcs', {
27+
'@google/genai': {
28+
GoogleGenAI: GoogleGenAI_Mock,
29+
},
30+
});
31+
32+
async function getGcsOutputUri() {
33+
return {
34+
uri: 'gs://mock/output',
35+
async cleanup() {},
36+
};
37+
}
38+
39+
describe('batchpredict-with-gcs', () => {
40+
it('should return the batch job state', async function () {
41+
this.timeout(500000);
42+
this.retries(4);
43+
await delay(this.test);
44+
const gcsOutput = await getGcsOutputUri();
45+
try {
46+
const output = await sample.runBatchPredictionJob(
47+
gcsOutput.uri,
48+
projectId,
49+
location
50+
);
51+
assert.notEqual(output, undefined);
52+
} finally {
53+
await gcsOutput.cleanup();
54+
}
55+
});
56+
});

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

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
// Copyright 2025 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// https://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
'use strict';
16+
17+
const {assert} = require('chai');
18+
const {describe, it} = require('mocha');
19+
20+
const projectId = process.env.CAIP_PROJECT_ID;
21+
const location = 'us-central1';
22+
const {delay} = require('../../test/util');
23+
24+
const proxyquire = require('proxyquire');
25+
const {GoogleGenAI_Mock} = require('./batchprediction-utils');
26+
27+
const sample = proxyquire('../batchpredict-with-bq', {
28+
'@google/genai': {
29+
GoogleGenAI: GoogleGenAI_Mock,
30+
},
31+
});
32+
33+
async function getBqOutputUri() {
34+
return {
35+
uri: 'gs://mock/output',
36+
async cleanup() {},
37+
};
38+
}
39+
40+
describe('batchpredict-with-bq', () => {
41+
it('should return the batch job state', async function () {
42+
this.timeout(500000);
43+
this.retries(4);
44+
await delay(this.test);
45+
const bqOutput = await getBqOutputUri();
46+
try {
47+
const output = await sample.runBatchPredictionJob(
48+
bqOutput.uri,
49+
projectId,
50+
location
51+
);
52+
assert.notEqual(output, undefined);
53+
} finally {
54+
await bqOutput.cleanup();
55+
}
56+
});
57+
});

0 commit comments

Comments
 (0)