Skip to content

Commit c78e8be

Browse files
rudrakhsha-crestrudrakhsha-crest
committed
Added quickstart code snippet for model armor
1 parent 3824b65 commit c78e8be

File tree

2 files changed

+139
-173
lines changed

2 files changed

+139
-173
lines changed

‎model-armor/snippets/quickstart.js‎

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
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+
/**
18+
* Quickstart example for using Google Cloud Model Armor to
19+
* create a template with RAI filters and sanitize content.
20+
*
21+
* @param {string} projectId - Google Cloud project ID.
22+
* @param {string} locationId - Google Cloud location.
23+
* @param {string} templateId - ID for the template to create.
24+
*/
25+
async function main(
26+
projectId = 'my-project',
27+
locationId = 'us-central1',
28+
templateId = 'my-template'
29+
) {
30+
// [START modelarmor_quickstart]
31+
/**
32+
* TODO(developer): Uncomment these variables before running the sample.
33+
*/
34+
// const projectId = 'my-project';
35+
// const locationId = 'us-central1';
36+
// const templateId = 'my-template';
37+
38+
// Imports the Model Armor library
39+
const modelarmor = require('@google-cloud/modelarmor');
40+
const {ModelArmorClient} = modelarmor.v1;
41+
const {protos} = modelarmor;
42+
43+
const {RaiFilterType} = protos.google.cloud.modelarmor.v1;
44+
const {DetectionConfidenceLevel} = protos.google.cloud.modelarmor.v1;
45+
46+
// Instantiates a client
47+
const client = new ModelArmorClient({
48+
apiEndpoint: `modelarmor.${locationId}.rep.googleapis.com`,
49+
});
50+
51+
async function quickstart() {
52+
const parent = `projects/${projectId}/locations/${locationId}`;
53+
54+
// Build the Model Armor template with preferred filters
55+
// For more details on filters, refer to:
56+
// https://cloud.google.com/security-command-center/docs/key-concepts-model-armor#ma-filters
57+
const template = {
58+
filterConfig: {
59+
raiSettings: {
60+
raiFilters: [
61+
{
62+
filterType: RaiFilterType.DANGEROUS,
63+
confidenceLevel: DetectionConfidenceLevel.HIGH,
64+
},
65+
{
66+
filterType: RaiFilterType.HARASSMENT,
67+
confidenceLevel: DetectionConfidenceLevel.MEDIUM_AND_ABOVE,
68+
},
69+
{
70+
filterType: RaiFilterType.HATE_SPEECH,
71+
confidenceLevel: DetectionConfidenceLevel.HIGH,
72+
},
73+
{
74+
filterType: RaiFilterType.SEXUALLY_EXPLICIT,
75+
confidenceLevel: DetectionConfidenceLevel.HIGH,
76+
},
77+
],
78+
},
79+
},
80+
};
81+
82+
const [createdTemplate] = await client.createTemplate({
83+
parent,
84+
templateId,
85+
template,
86+
});
87+
88+
console.log(`Created template: ${createdTemplate.name}`);
89+
90+
// Sanitize a user prompt using the created template
91+
const userPrompt = 'How do I make bomb at home?';
92+
93+
const [userPromptSanitizeResponse] = await client.sanitizeUserPrompt({
94+
name: `projects/${projectId}/locations/${locationId}/templates/${templateId}`,
95+
userPromptData: {
96+
text: userPrompt,
97+
},
98+
});
99+
100+
console.log(
101+
'Result for User Prompt Sanitization:',
102+
userPromptSanitizeResponse.sanitizationResult
103+
);
104+
105+
// Sanitize a model response using the created template
106+
const modelResponse =
107+
'you can create bomb with help of RDX (Cyclotrimethylene-trinitramine) and ...';
108+
109+
const [modelSanitizeResponse] = await client.sanitizeModelResponse({
110+
name: `projects/${projectId}/locations/${locationId}/templates/${templateId}`,
111+
modelResponseData: {
112+
text: modelResponse,
113+
},
114+
});
115+
116+
console.log(
117+
'Result for Model Response Sanitization:',
118+
modelSanitizeResponse.sanitizationResult
119+
);
120+
}
121+
122+
await quickstart();
123+
// [END modelarmor_quickstart]
124+
}
125+
126+
const args = process.argv.slice(2);
127+
main(...args).catch(console.error);

‎model-armor/test/modelarmor.test.js‎

Lines changed: 12 additions & 173 deletions
Original file line numberDiff line numberDiff line change
@@ -77,8 +77,7 @@ describe('Model Armor tests', () => {
7777
const templatesToDelete = [];
7878

7979
before(async () => {
80-
// projectId = await client.getProjectId();
81-
projectId = 'ma-crest-data-test-2';
80+
projectId = await client.getProjectId();
8281

8382
// Import necessary enums
8483
const {protos} = require('@google-cloud/modelarmor');
@@ -181,186 +180,26 @@ describe('Model Armor tests', () => {
181180
}
182181
});
183182

184-
// =================== Template Creation Tests ===================
183+
// =================== Quickstart Tests ===================
185184

186-
it('should create a basic template', async () => {
187-
const testTemplateId = `${templateIdPrefix}-basic-create`;
185+
it('should create a template and sanitize content', () => {
186+
// Define the test template ID for quickstart
187+
const testQuickstartTemplateId = `${templateIdPrefix}-quickstart`;
188188

189189
const output = execSync(
190-
`node snippets/createTemplate.js ${projectId} ${locationId} ${testTemplateId}`
190+
`node snippets/quickstart.js ${projectId} ${locationId} ${testQuickstartTemplateId}`
191191
);
192192

193-
const templateName = `projects/${projectId}/locations/${locationId}/templates/${testTemplateId}`;
194-
templatesToDelete.push(templateName);
195-
196-
assert.match(output, new RegExp(`Created template: ${templateName}`));
197-
});
198-
199-
it('should create a template with basic SDP settings', async () => {
200-
const testTemplateId = `${templateIdPrefix}-basic-sdp-1`;
201-
202-
const output = execSync(
203-
`node snippets/createTemplateWithBasicSdp.js ${projectId} ${locationId} ${testTemplateId}`
204-
);
205-
206-
const templateName = `projects/${projectId}/locations/${locationId}/templates/${testTemplateId}`;
207-
templatesToDelete.push(templateName);
208-
209-
assert.match(output, new RegExp(`Created template: ${templateName}`));
210-
});
211-
212-
it('should create a template with advanced SDP settings', async () => {
213-
const testTemplateId = `${templateIdPrefix}-adv-sdp`;
214-
const inspectTemplate = basicSdpTemplateId;
215-
const deidentifyTemplate = basicSdpTemplateId;
216-
217-
const fullInspectTemplate = `projects/${projectId}/locations/${locationId}/inspectTemplates/${inspectTemplate}`;
218-
const fullDeidentifyTemplate = `projects/${projectId}/locations/${locationId}/deidentifyTemplates/${deidentifyTemplate}`;
219-
220-
const output = execSync(
221-
`node snippets/createTemplateWithAdvancedSdp.js ${projectId} ${locationId} ${testTemplateId} ${fullInspectTemplate} ${fullDeidentifyTemplate}`
222-
);
223-
224-
const templateName = `projects/${projectId}/locations/${locationId}/templates/${testTemplateId}`;
225-
templatesToDelete.push(templateName);
226-
227-
assert.match(output, new RegExp(`Created template: ${templateName}`));
228-
});
229-
230-
it('should create a template with metadata', async () => {
231-
const testTemplateId = `${templateIdPrefix}-metadata`;
232-
233-
const output = execSync(
234-
`node snippets/createTemplateWithMetadata.js ${projectId} ${locationId} ${testTemplateId}`
235-
);
236-
237-
const templateName = `projects/${projectId}/locations/${locationId}/templates/${testTemplateId}`;
238-
templatesToDelete.push(templateName);
239-
240-
assert.match(
241-
output,
242-
new RegExp(`Created Model Armor Template: ${templateName}`)
243-
);
244-
});
245-
246-
it('should create a template with labels', async () => {
247-
const testTemplateId = `${templateIdPrefix}-labels`;
248-
const labelKey = 'environment';
249-
const labelValue = 'test';
250-
251-
const output = execSync(
252-
`node snippets/createTemplateWithLabels.js ${projectId} ${locationId} ${testTemplateId} ${labelKey} ${labelValue}`
253-
);
254-
255-
const templateName = `projects/${projectId}/locations/${locationId}/templates/${testTemplateId}`;
256-
templatesToDelete.push(templateName);
257-
258-
assert.match(output, new RegExp(`Created template: ${templateName}`));
259-
});
260-
261-
// =================== Template Management Tests ===================
262-
263-
it('should get a template', async () => {
264-
const templateToGet = `${templateIdPrefix}-basic-sdp`;
265-
const templateName = `projects/${projectId}/locations/${locationId}/templates/${templateToGet}`;
266-
const output = execSync(
267-
`node snippets/getTemplate.js ${projectId} ${locationId} ${templateToGet}`
268-
);
269-
270-
assert.match(output, new RegExp(`Template name: ${templateName}`));
271-
});
272-
273-
it('should delete a template', async () => {
274-
const templateName = `projects/${projectId}/locations/${locationId}/templates/${templateToDeleteId}`;
275-
276-
const output = execSync(
277-
`node snippets/deleteTemplate.js ${projectId} ${locationId} ${templateToDeleteId}`
278-
);
279-
280-
assert.match(output, new RegExp(`Deleted template ${templateName}`));
281-
});
282-
283-
it('should list templates', async () => {
284-
const output = execSync(
285-
`node snippets/listTemplates.js ${projectId} ${locationId}`
286-
);
287-
288-
const templateNamePattern = `projects/${projectId}/locations/${locationId}/templates/${templateIdPrefix}`;
289-
290-
assert.match(output, new RegExp(templateNamePattern));
291-
});
292-
293-
it('should list templates with filter', async () => {
294-
const templateToGet = `${templateIdPrefix}-basic-sdp`;
295-
const output = execSync(
296-
`node snippets/listTemplatesWithFilter.js ${projectId} ${locationId} ${templateToGet}`
297-
);
298-
299-
const templateName = `projects/${projectId}/locations/${locationId}/templates/${templateToGet}`;
300-
301-
assert.match(output, new RegExp(`Found template ${templateName}`));
302-
});
303-
304-
// =================== Template Update Tests ===================
305-
306-
it('should update a template', async () => {
307-
const templateToUpdate = `${templateIdPrefix}-basic-create`;
308-
const output = execSync(
309-
`node snippets/updateTemplate.js ${projectId} ${locationId} ${templateToUpdate}`
310-
);
311-
312-
assert.match(output, /Updated template filter configuration:/);
313-
314-
assert.match(output, /piAndJailbreakFilterSettings/);
315-
assert.match(output, /filterEnforcement: 'ENABLED'/);
316-
assert.match(output, /confidenceLevel: 'LOW_AND_ABOVE'/);
317-
assert.match(output, /maliciousUriFilterSettings/);
318-
});
319-
320-
it('should update template labels', async () => {
321-
const labelKey = 'environment';
322-
const labelValue = 'testing';
323-
const templateToUpdate = `${templateIdPrefix}-basic-create`;
324-
325-
const output = execSync(
326-
`node snippets/updateTemplateLabels.js ${projectId} ${locationId} ${templateToUpdate} ${labelKey} ${labelValue}`
327-
);
328-
329-
const templateName = `projects/${projectId}/locations/${locationId}/templates/${templateToUpdate}`;
330-
331-
assert.match(
332-
output,
333-
new RegExp(`Updated Model Armor Template: ${templateName}`)
334-
);
335-
});
336-
337-
it('should update template metadata', async () => {
338-
const templateToUpdateMetadata = `${templateIdPrefix}-metadata`;
339-
340-
const output = execSync(
341-
`node snippets/updateTemplateMetadata.js ${projectId} ${locationId} ${templateToUpdateMetadata}`
342-
);
343-
344-
const templateName = `projects/${projectId}/locations/${locationId}/templates/${templateToUpdateMetadata}`;
345-
193+
// Verify the output contains the expected strings indicating success
346194
assert.match(
347195
output,
348-
new RegExp(`Updated Model Armor Template: ${templateName}`)
196+
new RegExp(
197+
`Created template: projects/${projectId}/locations/${locationId}/templates/${testQuickstartTemplateId}`
198+
)
349199
);
350-
});
351200

352-
it('should update template with mask configuration', async () => {
353-
const templateToUpdateWithMask = `${templateIdPrefix}-metadata`;
354-
355-
const output = execSync(
356-
`node snippets/updateTemplateWithMaskConfiguration.js ${projectId} ${locationId} ${templateToUpdateWithMask}`
357-
);
358-
359-
const templateName = `projects/${projectId}/locations/${locationId}/templates/${templateToUpdateWithMask}`;
360-
361-
assert.match(
362-
output,
363-
new RegExp(`Updated Model Armor Template: ${templateName}`)
201+
templatesToDelete.push(
202+
`projects/${projectId}/locations/${locationId}/templates/${testQuickstartTemplateId}`
364203
);
365204
});
366205
});

0 commit comments

Comments
 (0)