diff --git a/ai-platform/snippets/predict-code-generation-function.js b/ai-platform/snippets/predict-code-generation-function.js new file mode 100644 index 0000000000..823b0d3de0 --- /dev/null +++ b/ai-platform/snippets/predict-code-generation-function.js @@ -0,0 +1,86 @@ +/* + * Copyright 2023 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'; + +async function main(project, location = 'us-central1') { + // [START aiplatform_sdk_code_generation_function] + /** + * TODO(developer): Uncomment these variables before running the sample.\ + * (Not necessary if passing values as arguments) + */ + // const project = 'YOUR_PROJECT_ID'; + // const location = 'YOUR_PROJECT_LOCATION'; + const aiplatform = require('@google-cloud/aiplatform'); + + // Imports the Google Cloud Prediction service client + const {PredictionServiceClient} = aiplatform.v1; + + // Import the helper module for converting arbitrary protobuf.Value objects. + const {helpers} = aiplatform; + + // Specifies the location of the api endpoint + const clientOptions = { + apiEndpoint: 'us-central1-aiplatform.googleapis.com', + }; + const publisher = 'google'; + const model = 'code-bison@001'; + + // Instantiates a client + const predictionServiceClient = new PredictionServiceClient(clientOptions); + + async function callPredict() { + // Configure the parent resource + const endpoint = `projects/${project}/locations/${location}/publishers/${publisher}/models/${model}`; + + const prompt = { + prefix: 'Write a function that checks if a year is a leap year.', + }; + const instanceValue = helpers.toValue(prompt); + const instances = [instanceValue]; + + const parameter = { + temperature: 0.2, + maxOutputTokens: 256, + }; + const parameters = helpers.toValue(parameter); + + const request = { + endpoint, + instances, + parameters, + }; + + // Predict request + const [response] = await predictionServiceClient.predict(request); + console.log('Get code generation response'); + const predictions = response.predictions; + console.log('\tPredictions :'); + for (const prediction of predictions) { + console.log(`\t\tPrediction : ${JSON.stringify(prediction)}`); + } + } + + callPredict(); + // [END aiplatform_sdk_code_generation_function] +} + +process.on('unhandledRejection', err => { + console.error(err.message); + process.exitCode = 1; +}); + +main(...process.argv.slice(2)); diff --git a/ai-platform/snippets/test/predict-code-generation-function.test.js b/ai-platform/snippets/test/predict-code-generation-function.test.js new file mode 100644 index 0000000000..85b6a983cf --- /dev/null +++ b/ai-platform/snippets/test/predict-code-generation-function.test.js @@ -0,0 +1,40 @@ +/* + * Copyright 2023 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 path = require('path'); +const {assert} = require('chai'); +const {describe, it} = require('mocha'); + +const cp = require('child_process'); +const execSync = cmd => cp.execSync(cmd, {encoding: 'utf-8'}); +const cwd = path.join(__dirname, '..'); + +const project = process.env.CAIP_PROJECT_ID; +const location = 'us-central1'; + +describe('AI platform predict code generation', () => { + it('should make predictions using a large language model', async () => { + const stdout = execSync( + `node ./predict-code-generation-function.js ${project} ${location}`, + { + cwd, + } + ); + assert.match(stdout, /Get code generation response/); + }); +});