From 3ab88f4c097dbbca9a18b385d3848f6a7b0d7f59 Mon Sep 17 00:00:00 2001
From: Dennis Traub <traubd@amazon.com>
Date: Fri, 15 Mar 2024 14:48:15 +0100
Subject: [PATCH] Refactor tools

---
 .../bedrock-runtime/tools/user_input.js       | 157 +++++++++---------
 1 file changed, 80 insertions(+), 77 deletions(-)

diff --git a/javascriptv3/example_code/bedrock-runtime/tools/user_input.js b/javascriptv3/example_code/bedrock-runtime/tools/user_input.js
index 4122bf8bea3..691b483680b 100644
--- a/javascriptv3/example_code/bedrock-runtime/tools/user_input.js
+++ b/javascriptv3/example_code/bedrock-runtime/tools/user_input.js
@@ -3,96 +3,99 @@
 
 import * as readline from "readline";
 
-function askAgain(askQuestion, newQuestion) {
-    readline.moveCursor(process.stdout, 0, -1);
-    readline.clearLine(process.stdout, 0);
-    readline.cursorTo(process.stdout, 0);
-    askQuestion(newQuestion);
-}
-
-function createInterface() {
-    return readline.createInterface({
-        input: process.stdin,
-        output: process.stdout
-    });
-}
 /**
  * @typedef {Object} FoundationModel
  * @property {string} modelName - The name of the model
  */
 
-/**
- * @param {FoundationModel[]} models - An array of models to choose from
- * @returns {Promise<FoundationModel>} - A Promise that resolves with the selected model
- */
-export const selectModel = (models) => {
-    return new Promise(resolve => {
-        const rl = createInterface();
+const rl = readline.createInterface({
+  input: process.stdin,
+  output: process.stdout,
+});
+
+const ask = (question, validate, onValid, repeat) => {
+  if (repeat) {
+    // Overwrite previous line
+    readline.moveCursor(process.stdout, 0, -1);
+    readline.clearLine(process.stdout, 0);
+    readline.cursorTo(process.stdout, 0);
+  }
 
-        const printOptions = () => {
-            models.forEach((model, index) => {
-                console.log(`${index + 1}. ${model.modelName}`);
-            });
-        };
+  rl.question(question, (answer) => {
+    if (validate(answer)) {
+      onValid(answer);
+    } else {
+      ask(question, validate, onValid, true);
+    }
+  });
+};
 
-        const askForModel = (question) => {
-            rl.question(question, answer => {
-                if (answer === "q") {
-                    rl.close();
-                    resolve(null);
-                }
-                else {
-                    const selectedIndex = parseInt(answer, 10) - 1;
-                    if (selectedIndex >= 0 && selectedIndex < models.length) {
-                        rl.close();
-                        resolve(models[selectedIndex]);
-                    } else {
-                        askAgain(askForModel, "Invalid input. Please enter a valid number (q to quit): ");
-                    }
-                }
-            });
-        };
+const select = (
+  /** @type {string[]} options */ options,
+  text,
+  validate,
+  onValid,
+) => {
+  // Print the options
+  options.forEach((option, index) => {
+    console.log(`${index + 1}. ${option}`);
+  });
 
-        printOptions();
-        askForModel("Select a model: (q to quit): ");
-    });
+  ask(text, validate, onValid, false);
 };
 
-export const askForPrompt = () => {
-    return new Promise(resolve => {
-        const rl = createInterface();
+export const selectModel = (/** @type {FoundationModel[]} */ models) => {
+  return new Promise((resolve) => {
+    const validate = (answer) => {
+      if (answer === "q") return true;
+      else {
+        const selectedIndex = parseInt(answer, 10) - 1;
+        return selectedIndex >= 0 && selectedIndex < models.length;
+      }
+    };
 
-        const askForPrompt = (question) => {
-            rl.question(question, answer => {
-                if (answer.trim() === "") {
-                    askAgain(askForPrompt, "Invalid input. Please enter a prompt: ");
-                } else {
-                    rl.close();
-                    resolve(answer);
-                }
-            });
-        };
-        askForPrompt("Now, enter your prompt: ");
-    });
+    const onValid = (answer) => {
+      if (answer === "q") {
+        rl.close();
+        resolve(null);
+      } else {
+        resolve(models[parseInt(answer, 10) - 1]);
+      }
+    };
+
+    const text = "Select a model number (q to quit): ";
+    select(
+      models.map((m) => m.modelName),
+      text,
+      validate,
+      onValid,
+    );
+  });
 };
 
-export const askForChoice = () => {
-    return new Promise((resolve) => {
-        const rl = createInterface();
+export const selectNextStep = (/** @type {string} */ modelName) => {
+  return new Promise((resolve) => {
+    const options = [`Prompt ${modelName} again`, "Select another model"];
+    const text = "Choose your next step (q to quit): ";
+
+    const validate = (answer) => ["1", "2", "q"].includes(answer);
+    const onValid = (answer) => {
+      if (answer === "q") {
+        rl.close();
+        resolve(null);
+      } else {
+        resolve(answer);
+      }
+    };
 
-        const askForChoice = (question) => {
-            rl.question(question, (answer) => {
-                if (["1", "2", "q"].includes(answer)) {
-                    rl.close();
-                    resolve(answer);
-                } else {
-                    askAgain(askForChoice, "Invalid input. Please enter 1, 2, or q: ");
-                }
-            });
-        };
+    select(options, text, validate, onValid);
+  });
+};
 
-        askForChoice(
-            "Enter 1 for a new prompt to the same model, 2 for a different model, or q to quit: "
-        );
-    });
+export const askForPrompt = () => {
+  return new Promise((resolve) => {
+    const validate = (/** @type {string} */ answer) => answer.trim() !== "";
+    const onValid = (answer) => resolve(answer);
+    ask("Now, enter your prompt: ", validate, onValid);
+  });
 };