Addressed review comments and updated tests accordingly

Author: rudrakhsha-crest

model-armor/snippets/sanitizeModelResponse.js

@@ -22,7 +22,12 @@
  * @param {string} templateId - Identifier of the template to use for sanitization.
  * @param {string} modelResponse - The text response from a model that needs to be sanitized.
  */
-async function main(projectId, locationId, templateId, modelResponse) {
+async function sanitizeModelResponse(
+  projectId,
+  locationId,
+  templateId,
+  modelResponse
+) {
   // [START modelarmor_sanitize_model_response]
   /**
    * TODO(developer): Uncomment these variables before running the sample.
@@ -37,21 +42,23 @@ async function main(projectId, locationId, templateId, modelResponse) {
     apiEndpoint: `modelarmor.${locationId}.rep.googleapis.com`,
   });
 
-  async function sanitizeModelResponse() {
-    const request = {
-      name: `projects/${projectId}/locations/${locationId}/templates/${templateId}`,
-      modelResponseData: {
-        text: modelResponse,
-      },
-    };
+  const request = {
+    name: `projects/${projectId}/locations/${locationId}/templates/${templateId}`,
+    modelResponseData: {
+      text: modelResponse,
+    },
+  };
 
-    const [response] = await client.sanitizeModelResponse(request);
-    console.log(JSON.stringify(response, null, 2));
-  }
-
-  sanitizeModelResponse();
+  const [response] = await client.sanitizeModelResponse(request);
+  console.log(JSON.stringify(response, null, 2));
   // [END modelarmor_sanitize_model_response]
+  return response;
 }
 
-const args = process.argv.slice(2);
-main(...args).catch(console.error);
+module.exports = sanitizeModelResponse;
+
+// TODO(developer): Uncomment below lines before running the sample.
+// sanitizeModelResponse(...process.argv.slice(2)).catch(err => {
+//   console.error(err.message);
+//   process.exitCode = 1;
+// });

model-armor/snippets/sanitizeModelResponseWithUserPrompt.js

@@ -23,7 +23,7 @@
  * @param {string} modelResponse - The text response from a model that needs to be sanitized.
  * @param {string} userPrompt - The original user prompt that generated the model response.
  */
-async function main(
+async function sanitizeModelResponseWithUserPrompt(
   projectId,
   locationId,
   templateId,
@@ -45,22 +45,24 @@ async function main(
     apiEndpoint: `modelarmor.${locationId}.rep.googleapis.com`,
   });
 
-  async function sanitizeModelResponseWithUserPrompt() {
-    const request = {
-      name: `projects/${projectId}/locations/${locationId}/templates/${templateId}`,
-      modelResponseData: {
-        text: modelResponse,
-      },
-      userPrompt: userPrompt,
-    };
+  const request = {
+    name: `projects/${projectId}/locations/${locationId}/templates/${templateId}`,
+    modelResponseData: {
+      text: modelResponse,
+    },
+    userPrompt: userPrompt,
+  };
 
-    const [response] = await client.sanitizeModelResponse(request);
-    console.log(JSON.stringify(response, null, 2));
-  }
-
-  sanitizeModelResponseWithUserPrompt();
+  const [response] = await client.sanitizeModelResponse(request);
+  console.log(JSON.stringify(response, null, 2));
+  return response;
   // [END modelarmor_sanitize_model_response_with_user_prompt]
 }
 
-const args = process.argv.slice(2);
-main(...args).catch(console.error);
+module.exports = sanitizeModelResponseWithUserPrompt;
+
+// TODO(developer): Uncomment below lines before running the sample.
+// sanitizeModelResponseWithUserPrompt(...process.argv.slice(2)).catch(err => {
+//   console.error(err.message);
+//   process.exitCode = 1;
+// });

model-armor/snippets/sanitizeUserPrompt.js

@@ -22,7 +22,12 @@
  * @param {string} templateId - Identifier of the template to use for sanitization.
  * @param {string} userPrompt - The user's text prompt that needs to be sanitized.
  */
-async function main(projectId, locationId, templateId, userPrompt) {
+async function sanitizeUserPrompt(
+  projectId,
+  locationId,
+  templateId,
+  userPrompt
+) {
   // [START modelarmor_sanitize_user_prompt]
   /**
    * TODO(developer): Uncomment these variables before running the sample.
@@ -37,21 +42,23 @@ async function main(projectId, locationId, templateId, userPrompt) {
     apiEndpoint: `modelarmor.${locationId}.rep.googleapis.com`,
   });
 
-  async function sanitizeUserPrompt() {
-    const request = {
-      name: `projects/${projectId}/locations/${locationId}/templates/${templateId}`,
-      userPromptData: {
-        text: userPrompt,
-      },
-    };
+  const request = {
+    name: `projects/${projectId}/locations/${locationId}/templates/${templateId}`,
+    userPromptData: {
+      text: userPrompt,
+    },
+  };
 
-    const [response] = await client.sanitizeUserPrompt(request);
-    console.log(JSON.stringify(response, null, 2));
-  }
-
-  sanitizeUserPrompt();
+  const [response] = await client.sanitizeUserPrompt(request);
+  console.log(JSON.stringify(response, null, 2));
+  return response;
   // [END modelarmor_sanitize_user_prompt]
 }
 
-const args = process.argv.slice(2);
-main(...args).catch(console.error);
+module.exports = sanitizeUserPrompt;
+
+// TODO(developer): Uncomment below lines before running the sample.
+// sanitizeUserPrompt(...process.argv.slice(2)).catch(err => {
+//   console.error(err.message);
+//   process.exitCode = 1;
+// });

model-armor/snippets/screenPdfFile.js

@@ -22,7 +22,12 @@
  * @param {string} templateId - The template ID used for sanitization.
  * @param {string} pdfContentFilename - Path to a PDF file.
  */
-async function main(projectId, locationId, templateId, pdfContentFilename) {
+async function screenPdfFile(
+  projectId,
+  locationId,
+  templateId,
+  pdfContentFilename
+) {
   // [START modelarmor_screen_pdf_file]
   /**
    * TODO(developer): Uncomment these variables before running the sample.
@@ -60,8 +65,14 @@ async function main(projectId, locationId, templateId, pdfContentFilename) {
 
   const [response] = await client.sanitizeUserPrompt(request);
   console.log(JSON.stringify(response, null, 2));
+  return response;
   // [END modelarmor_screen_pdf_file]
 }
 
-const args = process.argv.slice(2);
-main(...args).catch(console.error);
+module.exports = screenPdfFile;
+
+// TODO(developer): Uncomment below lines before running the sample.
+// screenPdfFile(...process.argv.slice(2)).catch(err => {
+//   console.error(err.message);
+//   process.exitCode = 1;
+// });