Clean up the chatbot module (postgresml#1274)

open-prophetdb · Jan 5, 2024 · 4b51231 · 4b51231
1 parent abb9533
commit 4b51231
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 23 deletions.
diff --git a/pgml-dashboard/src/components/chatbot/chatbot_controller.js b/pgml-dashboard/src/components/chatbot/chatbot_controller.js
@@ -174,13 +174,15 @@ export default class extends Controller {
     this.handleKnowledgeBaseChange(); // This will set our initial knowledge base
     this.handleBrainChange(); // This will set our initial brain
     this.handleResize();
+    this.openConnection();
+    this.getHistory();
+  }
 
+  openConnection() {
     const url = ((window.location.protocol === "https:") ? "wss://" : "ws://") + window.location.hostname + (((window.location.port != 80) && (window.location.port != 443)) ? ":" + window.location.port : "") + window.location.pathname + "/get-answer";
     this.socket = new WebSocket(url);
     this.socket.onmessage = (message) => {
       let result = JSON.parse(message.data);
-      console.log(result);
-
       if (result.error) {
         this.showChatbotAlert("Error", "Error getting chatbot answer");
         console.log(result.error);
@@ -201,7 +203,6 @@ export default class extends Controller {
     this.socket.onclose = () => {
       window.setTimeout(() => this.openConnection(), 500);
     };
-    this.getHistory();
   }
 
   async clearHistory() {

diff --git a/pgml-dashboard/src/components/chatbot/mod.rs b/pgml-dashboard/src/components/chatbot/mod.rs
@@ -72,31 +72,31 @@ impl KnowledgeBaseWithLogo {
     }
 }
 
-const CHATBOT_BRAINS: [ChatbotBrain; 4] = [
-    ChatbotBrain::new(
-        "teknium/OpenHermes-2.5-Mistral-7B",
-        "OpenHermes",
-        "teknium/OpenHermes-2.5-Mistral-7B",
-        "/dashboard/static/images/logos/openhermes.webp",
-    ),
-    ChatbotBrain::new(
-        "Gryphe/MythoMax-L2-13b",
-        "MythoMax",
-        "Gryphe/MythoMax-L2-13b",
-        "/dashboard/static/images/logos/mythomax.webp",
-    ),
+const CHATBOT_BRAINS: [ChatbotBrain; 1] = [
+    // ChatbotBrain::new(
+    //     "teknium/OpenHermes-2.5-Mistral-7B",
+    //     "OpenHermes",
+    //     "teknium/OpenHermes-2.5-Mistral-7B",
+    //     "/dashboard/static/images/logos/openhermes.webp",
+    // ),
+    // ChatbotBrain::new(
+    //     "Gryphe/MythoMax-L2-13b",
+    //     "MythoMax",
+    //     "Gryphe/MythoMax-L2-13b",
+    //     "/dashboard/static/images/logos/mythomax.webp",
+    // ),
     ChatbotBrain::new(
         "openai",
         "OpenAI",
         "ChatGPT",
         "/dashboard/static/images/logos/openai.webp",
     ),
-    ChatbotBrain::new(
-        "berkeley-nest/Starling-LM-7B-alpha",
-        "Starling",
-        "berkeley-nest/Starling-LM-7B-alpha",
-        "/dashboard/static/images/logos/starling.webp",
-    ),
+    // ChatbotBrain::new(
+    //     "berkeley-nest/Starling-LM-7B-alpha",
+    //     "Starling",
+    //     "berkeley-nest/Starling-LM-7B-alpha",
+    //     "/dashboard/static/images/logos/starling.webp",
+    // ),
 ];
 
 struct ChatbotBrain {
@@ -125,7 +125,7 @@ impl ChatbotBrain {
 #[derive(TemplateOnce)]
 #[template(path = "chatbot/template.html")]
 pub struct Chatbot {
-    brains: &'static [ChatbotBrain; 4],
+    brains: &'static [ChatbotBrain; 1],
     example_questions: &'static ExampleQuestions,
     knowledge_bases_with_logo: &'static [KnowledgeBaseWithLogo; 4],
 }