feat: wikipedia module

1 files changed, 16 insertions, 10 deletions

diff --git a/src/app/init.rs b/src/app/init.rs
index dd8d3b7..0930319 100644
--- a/src/app/init.rs
+++ b/src/app/init.rs
@@ -1,27 +1,30 @@
 use crate::app::llm::{Message, MessageType, LLM};
+use crate::app::modules::wikipedia::ask_wiki;
 use crate::helper::init::warn;
 use uuid::Uuid;
 use tokio::runtime::Builder;
 
 pub struct App {
     pub messages: Vec<Message>, // History of recorded message
-    conv_id: Uuid,
-    chat_llm: LLM,
-    resume_llm: LLM,
+    pub conv_id: Uuid, // ID for retrieving and saving the history of messag
+    categorize_llm: LLM,
+    chat_llm: LLM, // Configuration for the LLM that chat with you
+    resume_llm: LLM, // Configuration for the LLM that resume conversation
 }
 
 impl App {
     pub fn new() -> App {
-        let chat_llm: LLM = LLM::new("config/chat-LLM.json".to_string());
 
+        let categorize_llm = LLM::new("config/categorize-LLM.json");
         App {
             messages: vec![Message::new(
                 MessageType::SYSTEM,
-                chat_llm.system_prompt.clone(),
+                categorize_llm.system_prompt.clone(),
             )],
             conv_id: Uuid::new_v4(),
-            chat_llm,
-            resume_llm: LLM::new("config/resume-LLM.json".to_string()),
+            categorize_llm,
+            chat_llm: LLM::new("config/chat-LLM.json"),
+            resume_llm: LLM::new("config/resume-LLM.json"),
         }
     }
 
@@ -29,7 +32,7 @@ impl App {
         let message = Message::new(role, msg);
 
         let err = message.save_message(self.conv_id.to_string());
-        warn(err.is_err().to_string()); 
+        //warn(err.is_err().to_string()); 
 
         self.messages.push(message);
     }
@@ -39,19 +42,20 @@ impl App {
 
         let result = runtime.block_on(async {
             // Ask the LLM to categorise the request between (chat, code, wikipedia)
-            self.chat_llm.ask_format(&self.messages).await
+            self.categorize_llm.ask_tools(&self.messages).await
         });
 
         match result {
             Ok(msg) => {
                 let categorie = msg[0]["function"]["arguments"]["category"].clone();
-                self.ask(categorie.to_string().as_str());
+                self.ask(&categorie.to_string().replace("\"", ""));
             },
             Err(e) => self.append_message(e.to_string(), MessageType::ASSISTANT),
         }
     }
 
     fn ask(&mut self, mode: &str) {
+        warn(format!("Categorie: {}", mode));
         let runtime = Builder::new_current_thread()
             .enable_all()
             .build().unwrap();
@@ -59,6 +63,8 @@ impl App {
         let result = runtime.block_on(async {
             if mode == "resume" {
                 self.resume_llm.ask(&self.messages).await
+            } else if mode == "wikipedia" {
+                ask_wiki(&self.messages).await
             } else {
                 self.chat_llm.ask(&self.messages).await
             }