From e5419b01086aa218ac2d61342bd22a158317f958 Mon Sep 17 00:00:00 2001
From: Wizzard <rich@bandaholics.cash>
Date: Sat, 9 Mar 2024 04:50:30 -0500
Subject: [PATCH] Update kuzco-cli to allow choice of model on startup &
 include a timeout function

---
 cli/kuzco-cli.js | 27 ++++++++++++++++++++++-----
 1 file changed, 22 insertions(+), 5 deletions(-)

diff --git a/cli/kuzco-cli.js b/cli/kuzco-cli.js
index 6450e83..49a867d 100644
--- a/cli/kuzco-cli.js
+++ b/cli/kuzco-cli.js
@@ -52,9 +52,27 @@ const askQuestion = (query) => {
     }));
 };
 
+const fetchWithTimeout = (url, options, timeout = 3000) => {
+    const timeoutPromise = new Promise((_, reject) => {
+        const id = setTimeout(() => {
+            clearTimeout(id);
+            reject(new Error('Request timed out'));
+        }, timeout);
+    });
+
+    return Promise.race([
+        fetch(url, options),
+        timeoutPromise
+    ]);
+};
+
 async function main() {
     let messages = [];
 
+    console.log("Please choose a model: 1 for Mistral, 2 for Llama2");
+    const modelChoice = prompt('Enter your choice (1 or 2): ');
+    const model = modelChoice === '2' ? 'llama2' : 'mistral';
+
     while (true) {
         const user_input = await askQuestion("User: ");
         if (user_input.toLowerCase() === 'exit') {
@@ -63,7 +81,7 @@ async function main() {
         messages.push({ 'role': 'user', 'content': user_input + '\n' });
 
         try {
-            const response = await fetch(`${BASE_URL}/chat/completions`, {
+            const response = await fetchWithTimeout(`${BASE_URL}/chat/completions`, {
                 method: 'POST',
                 headers: {
                     'Authorization': `Bearer ${API_KEY}`,
@@ -71,17 +89,16 @@ async function main() {
                 },
                 body: JSON.stringify({
                     messages: messages,
-                    model: 'mistral',
-                    stream: false
+                    model: model
                 })
-            });
+            }, 15000);;
 
             if (!response.ok) {
                 throw new Error(`HTTP error! status: ${response.status}`);
             }
 
             const data = await response.json();
-            console.log(`\nKuzco (Mistral):\n\n${data.choices[0].message.content.trim()}\n`);
+            console.log(`\nKuzco (${model}):\n\n${data.choices[0].message.content.trim()}\n`);
         } catch (error) {
             console.error(`An error occurred: ${error.message}`);
         }