Added images and finalized models for chatbot

postgresml · SilasMarvin · Jan 3, 2024 · Dec 5, 2023 · Dec 8, 2023 · Dec 14, 2023
commit aa82b6dd7e18c77b3baea4327dfe6a6ada829626
diff --git a/pgml-dashboard/src/api/chatbot.rs b/pgml-dashboard/src/api/chatbot.rs
@@ -48,7 +48,7 @@ impl ChatRole {
             ChatRole::Bot => match brain {
                 ChatbotBrain::OpenAIGPT4
                 | ChatbotBrain::TekniumOpenHermes25Mistral7B
-                | ChatbotBrain::Openchat35 => "assistant",
+                | ChatbotBrain::Starling7b => "assistant",
                 ChatbotBrain::GrypheMythoMaxL213b => "model",
             },
             ChatRole::System => "system",
@@ -61,7 +61,7 @@ enum ChatbotBrain {
     OpenAIGPT4,
     TekniumOpenHermes25Mistral7B,
     GrypheMythoMaxL213b,
-    Openchat35,
+    Starling7b,
 }
 
 impl ChatbotBrain {
@@ -104,13 +104,12 @@ impl ChatbotBrain {
                 "revision": "main",
                 "device_map": "auto"
             }),
-            // Self::GrypheMythoMaxL213b => serde_json::json!("Gryphe/MythoMax-L2-13b"),
             Self::GrypheMythoMaxL213b => serde_json::json!({
                 "model": "TheBloke/MythoMax-L2-13B-GPTQ",
                 "revision": "main",
                 "device_map": "auto"
             }),
-            Self::Openchat35 => serde_json::json!({
+            Self::Starling7b => serde_json::json!({
                 "model": "TheBloke/Starling-LM-7B-alpha-GPTQ",
                 "revision": "main",
                 "device_map": "auto"
@@ -136,7 +135,7 @@ impl TryFrom<&str> for ChatbotBrain {
             "teknium/OpenHermes-2.5-Mistral-7B" => Ok(ChatbotBrain::TekniumOpenHermes25Mistral7B),
             "Gryphe/MythoMax-L2-13b" => Ok(ChatbotBrain::GrypheMythoMaxL213b),
             "openai" => Ok(ChatbotBrain::OpenAIGPT4),
-            "openchat/openchat_3.5" => Ok(ChatbotBrain::Openchat35),
+            "berkeley-nest/Starling-LM-7B-alpha" => Ok(ChatbotBrain::Starling7b),
             _ => Err(anyhow::anyhow!("Invalid brain id")),
         }
     }
@@ -148,7 +147,7 @@ impl From<ChatbotBrain> for &'static str {
             ChatbotBrain::TekniumOpenHermes25Mistral7B => "teknium/OpenHermes-2.5-Mistral-7B",
             ChatbotBrain::GrypheMythoMaxL213b => "Gryphe/MythoMax-L2-13b",
             ChatbotBrain::OpenAIGPT4 => "openai",
-            ChatbotBrain::Openchat35 => "openchat/openchat_3.5",
+            ChatbotBrain::Starling7b => "berkeley-nest/Starling-LM-7B-alpha",
         }
     }
 }
@@ -375,7 +374,7 @@ struct HistoryMessage {
     side: String,
     content: String,
     knowledge_base: String,
-    model: String,
+    brain: String,
 }
 
 #[get("/chatbot/get-history")]
@@ -448,7 +447,7 @@ async fn do_chatbot_get_history(user: &User, limit: usize) -> anyhow::Result<Vec
             Ok(HistoryMessage {
                 side,
                 content,
-                model: model.to_string(),
+                brain: model.to_string(),
                 knowledge_base: knowledge_base.to_string(),
             })
         })
@@ -525,7 +524,6 @@ async fn process_message(
         let collection = Collection::new(
             collection,
             Some(std::env::var("CHATBOT_DATABASE_URL").expect("CHATBOT_DATABASE_URL not set")),
-            // Some(std::env::var("CHATBOT_DATABASE_URL_TWO").expect("CHATBOT_DATABASE_URL not set")),
         );
         let context = collection
         .query()
@@ -541,7 +539,7 @@ async fn process_message(
         .join("\n");
         // let context = "".to_string();
 
-        let mut history_collection = Collection::new(
+        let history_collection = Collection::new(
             "ChatHistory",
             Some(std::env::var("CHATBOT_DATABASE_URL").expect("CHATBOT_DATABASE_URL not set")),
         );

diff --git a/pgml-dashboard/src/components/chatbot/chatbot_controller.js b/pgml-dashboard/src/components/chatbot/chatbot_controller.js
@@ -15,23 +15,23 @@ const LOADING_MESSAGE = `
 </div>
 `;
 
-const getBackgroundImageURLForSide = (side, knowledgeBase) => {
+const getBackgroundImageURLForSide = (side, brain) => {
   if (side == "user") {
     return "/dashboard/static/images/chatbot_user.webp";
   } else {
-    if (knowledgeBase == "postgresml") {
-      return "/dashboard/static/images/owl_gradient.svg";
-    } else if (knowledgeBase == "pytorch") {
-      return "/dashboard/static/images/logos/pytorch.svg";
-    } else if (knowledgeBase == "rust") {
-      return "/dashboard/static/images/logos/rust.svg";
-    } else if (knowledgeBase == "postgresql") {
-      return "/dashboard/static/images/logos/postgresql.svg";
+    if (brain == "teknium/OpenHermes-2.5-Mistral-7B") {
+      return "/dashboard/static/images/logos/openhermes.webp"
+    } else if (brain == "Gryphe/MythoMax-L2-13b") {
+      return "/dashboard/static/images/logos/mythomax.webp"
+    } else if (brain == "berkeley-nest/Starling-LM-7B-alpha") {
+      return "/dashboard/static/images/logos/starling.webp"    
+    } else if (brain == "openai") {
+      return "/dashboard/static/images/logos/openai.webp"
     }
   }
 };
 
-const createHistoryMessage = (message, knowledgeBase) => {
+const createHistoryMessage = (message) => {
   if (message.side == "system") {
     return `
       <div class="chatbot-knowledge-base-notice text-center p-3">${message.text}</div>
@@ -46,7 +46,7 @@ const createHistoryMessage = (message, knowledgeBase) => {
           <div class="rounded p-1 chatbot-message-avatar-wrapper">
             <div class="chatbot-message-avatar" style="background-image: url('${getBackgroundImageURLForSide(
               message.side,
-              knowledgeBase,
+              message.brain,
             )}')">
           </div>
         </div>
@@ -76,8 +76,8 @@ const brainIdToName = (brain) => {
     return "OpenHermes"
   } else if (brain == "Gryphe/MythoMax-L2-13b") {
     return "MythoMax"
-  } else if (brain == "openchat/openchat_3.5") {
-    return "OpenChat"    
+  } else if (brain == "berkeley-nest/Starling-LM-7B-alpha") {
+    return "Starling"    
   } else if (brain == "openai") {
     return "ChatGPT"
   }
@@ -92,10 +92,11 @@ const createKnowledgeBaseNotice = (knowledgeBase) => {
 };
 
 class Message {
-  constructor(id, side, text, is_partial=false) {
+  constructor(id, side, brain, text, is_partial=false) {
     this.id = id
-    this.text = text
     this.side = side
+    this.brain = brain
+    this.text = text
     this.is_partial = is_partial
   }
 
@@ -187,9 +188,9 @@ export default class extends Controller {
       } else {
         let message;
         if (result.partial_result) {
-          message = new Message(result.id, "bot", result.partial_result, true);
+          message = new Message(result.id, "bot", this.brain, result.partial_result, true);
         } else {
-          message = new Message(result.id, "bot", result.result, false);
+          message = new Message(result.id, "bot", this.brain, result.result);
         }
         this.messageHistory.add_message(message, this.messageIdToKnowledgeBaseId[message.id]);
         this.redrawChat();
@@ -216,7 +217,7 @@ export default class extends Controller {
       console.log("Error getting chat history", history.error)
     } else {
       for (const message of history.result) {
-        const newMessage = new Message(getRandomInt(), message.side, message.content, false);
+        const newMessage = new Message(getRandomInt(), message.side, message.brain, message.content, false);
         console.log(newMessage);
         this.messageHistory.add_message(newMessage, message.knowledge_base);
       }
@@ -231,10 +232,7 @@ export default class extends Controller {
       console.log("Drawing", message);
       this.chatHistory.insertAdjacentHTML(
         "beforeend",
-        createHistoryMessage(
-          message,
-          this.knowledgeBase,
-        ),
+        createHistoryMessage(message),
       );
     }
 
@@ -250,19 +248,16 @@ export default class extends Controller {
   }
 
   newUserQuestion(question) {
-    const message = new Message(getRandomInt(), "user", question);
+    const message = new Message(getRandomInt(), "user", this.brain, question);
     this.messageHistory.add_message(message, this.knowledgeBase);
     this.messageIdToKnowledgeBaseId[message.id] = this.knowledgeBase;
     this.hideExampleQuestions();
     this.redrawChat();
 
-    let loadingMessage = new Message("loading", "bot", LOADING_MESSAGE);
+    let loadingMessage = new Message("loading", "bot", this.brain, LOADING_MESSAGE);
     this.chatHistory.insertAdjacentHTML(
       "beforeend",
-      createHistoryMessage(
-        loadingMessage,
-        this.knowledgeBase,
-      ),
+      createHistoryMessage(loadingMessage),
     );
     this.chatHistory.scrollTop = this.chatHistory.scrollHeight;
 
@@ -331,7 +326,7 @@ export default class extends Controller {
     let knowledge_base = knowledgeBaseIdToName(this.knowledgeBase);
     let brain = brainIdToName(this.brain);
     let content = `Chatting with ${brain} about ${knowledge_base}`;
-    const newMessage = new Message(getRandomInt(), "system", content);
+    const newMessage = new Message(getRandomInt(), "system", this.brain, content);
     this.messageHistory.add_message(newMessage, this.knowledgeBase);
     this.redrawChat();
   }

diff --git a/pgml-dashboard/src/components/chatbot/mod.rs b/pgml-dashboard/src/components/chatbot/mod.rs
@@ -77,13 +77,13 @@ const CHATBOT_BRAINS: [ChatbotBrain; 4] = [
         "teknium/OpenHermes-2.5-Mistral-7B",
         "OpenHermes",
         "teknium/OpenHermes-2.5-Mistral-7B",
-        "/dashboard/static/images/owl_gradient.svg",
+        "/dashboard/static/images/logos/openhermes.webp",
     ),
     ChatbotBrain::new(
         "Gryphe/MythoMax-L2-13b",
         "MythoMax",
         "Gryphe/MythoMax-L2-13b",
-        "/dashboard/static/images/owl_gradient.svg",
+        "/dashboard/static/images/logos/mythomax.webp",
     ),
     ChatbotBrain::new(
         "openai",
@@ -92,10 +92,10 @@ const CHATBOT_BRAINS: [ChatbotBrain; 4] = [
         "/dashboard/static/images/logos/openai.webp",
     ),
     ChatbotBrain::new(
-        "openchat/openchat_3.5",
-        "OpenChat",
-        "openchat/openchat_3.5",
-        "/dashboard/static/images/logos/meta.webp",
+        "berkeley-nest/Starling-LM-7B-alpha",
+        "Starling",
+        "berkeley-nest/Starling-LM-7B-alpha",
+        "/dashboard/static/images/logos/starling.webp",
     ),
 ];
 

diff --git a/pgml-dashboard/static/images/logos/mythomax.webp b/pgml-dashboard/static/images/logos/mythomax.webp
diff --git a/pgml-dashboard/static/images/logos/openhermes.webp b/pgml-dashboard/static/images/logos/openhermes.webp
diff --git a/pgml-dashboard/static/images/logos/starling.webp b/pgml-dashboard/static/images/logos/starling.webp