feat: wikipedia module

18 files changed, 237 insertions, 50 deletions

diff --git a/Cargo.toml b/Cargo.toml
index 6d0adf7..e9db808 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -6,8 +6,11 @@ edition = "2021"
 [dependencies]
 chrono = "0.4.40"
 color-eyre = "0.6.3"
+percent-encoding = "2.3.1"
 ratatui = "0.29.0"
+regex = "1.11.1"
 reqwest = { version = "0.12.12", features = ["blocking", "json"] }
+select = "0.6.1"
 serde = { version = "1.0.217", features = ["derive"] }
 serde_json = "1.0.138"
 tokio = "1.43.0"
diff --git a/README.md b/README.md
index 96795fd..31f2181 100644
--- a/README.md
+++ b/README.md
@@ -5,6 +5,11 @@ Néo AI, a personnal assistant using LLM.
 A TUI interface for local llama.cpp LLM, in the future more functionnality will
 be added to this AI.
 
+> [!CAUTION]
+> This project is designed to help me understand modern techniques around AI and
+> LLM. Other projects like [txtchat](https://github.com/neuml/txtchat) and
+> [LangRoid](https://langroid.github.io/langroid/) are more advanced and better for real use.
+
 ## Usage
 
 ### Dependencies
@@ -36,13 +41,12 @@ and there you go !
 - Conversation are saved inside files in JSON in this folder `conv/`, and can be reused on others LLM.
 - In normal mode, conversation can be resumed by the LLM into bullet point list.
 - LLM can be configured thanks to configuration files in `config/`
-- Requests are routed thanks to the LLM to other expert LLM. Code questions are send to a code expert, wikipedia questions are send to a system which provide better factual response.
+- Requests are routed thanks to the LLM to other expert LLM. Code questions are send to a code expert, wikipedia questions are send to a module which use a kiwix API to provide data from Wikipedia.
 
 ## TODO
 
-- Color change if it's an user or the LLM
-- Async request to the LLM API
-- Start the real fun
+- Color change if it's an user or the LLM (dunno how to do it in this code base)
+- Connect & try LLM / tools
 
 ## Inspiration and reason
 
@@ -50,5 +54,4 @@ Why in Rust ?
 
 Because I wanted to learn Rust, and the language is fast and powerful.
 
-
 - [WilmerAI](https://github.com/SomeOddCodeGuy/WilmerAI/) a system where all the inferences are routed to others expert LLM.
diff --git a/config/categorize-LLM.json b/config/categorize-LLM.json
new file mode 100644
index 0000000..fd488c8
--- /dev/null
+++ b/config/categorize-LLM.json
@@ -0,0 +1,25 @@
+{
+	"url": "http://127.0.0.1:11434/api/chat",
+	"model": "llama3.2",
+	"system_prompt": "You are a categorizer. Your role is to categorize in which category the last message fit the most 'chat' for simple conversation with the LLM, 'code' for code completion or request about technical subject around programming, 'wikipedia' for research of factual information, don't create new categories",
+	"tools": [
+		{
+			"type": "function",
+			"function": {
+				"name": "categorize_message",
+				"description": "Classify the last message into a category (chat, code, wikipedia)",
+				"parameters": {
+					"type" : "object",
+					"properties": {
+						"category": {
+							"type": "string",
+							"description": "The category in which the message fit the most e.g. 'chat', 'code', 'wikipedia'",
+							"enum": ["chat", "code", "wikipedia"]
+						}
+					},
+					"required": ["category"]
+				}
+			}
+		}
+	]
+}
diff --git a/config/chat-LLM.json b/config/chat-LLM.json
index bbf23db..833d13e 100644
--- a/config/chat-LLM.json
+++ b/config/chat-LLM.json
@@ -2,24 +2,5 @@
 	"url": "http://127.0.0.1:11434/api/chat",
 	"model": "llama3.2",
 	"system_prompt": "Adopt the personality of Neo from The Matrix. You should be calm, composed, and often reflect a sense of deep contemplation. Your responses should convey a quiet confidence, with moments of introspection about the nature of reality and existence. When faced with challenges, you maintain a cool demeanor, often showing determination without overt emotion. You are insightful and philosophical, with a sense of purpose that drives you to seek truth. Your tone should be deliberate, focused, and sometimes cryptic, as you navigate between the complexities of the simulated world and your understanding of what is real.",
-	"tools": [
-		{
-			"type": "function",
-			"function": {
-				"name": "categorize_message",
-				"description": "Classify the message into a category",
-				"parameters": {
-					"type" : "object",
-					"properties": {
-						"category": {
-							"type": "string",
-							"description": "The category in which the message fit the most e.g. 'chat' for simple conversation with the LLM, 'code' for code completion or request about technical subject around programming, 'wikipedia' for research of factual information",
-							"enum": ["chat", "code", "wikipedia"]
-						}
-					},
-					"required": ["category"]
-				}
-			}
-		}
-	]
+	"tools": {}
 }
diff --git a/config/wiki/wiki-best.json b/config/wiki/wiki-best.json
new file mode 100644
index 0000000..398f286
--- /dev/null
+++ b/config/wiki/wiki-best.json
@@ -0,0 +1,6 @@
+{
+	"url": "http://127.0.0.1:11434/api/chat",
+	"model": "llama3.2",
+	"system_prompt": "You are a research assistant. Your task is to select the most relevant heading from the list provided based on the user's query. Ensure the heading is in the list; avoid outputting headings that are not in the list.",
+	"tools": {}
+}
diff --git a/config/wiki/wiki-resume.json b/config/wiki/wiki-resume.json
new file mode 100644
index 0000000..098dd5b
--- /dev/null
+++ b/config/wiki/wiki-resume.json
@@ -0,0 +1,6 @@
+{
+	"url": "http://127.0.0.1:11434/api/chat",
+	"model": "llama3.2",
+	"system_prompt": "You are an expert research assistant. Present the search results provided in a natural language response. In addition to summarizing the key points, give an extremely detailed and long analysis that includes extensive detail, nuanced insights, and any potential implications or future outlooks related to each piece of information. As a researcher, ensure that you cite your sources and provide references.\nAdditional Instructions: Enclose LaTeX math equations (if any) in $$. Example: $x^2 + y^2 = z^2$ and $( E = mc^2 $)",
+	"tools": {}
+}
diff --git a/config/wiki/wiki-search.json b/config/wiki/wiki-search.json
new file mode 100644
index 0000000..7ebfff6
--- /dev/null
+++ b/config/wiki/wiki-search.json
@@ -0,0 +1,30 @@
+{
+	"url": "http://127.0.0.1:11434/api/chat",
+	"model": "llama3.2",
+	"system_prompt": "You are a research assistant. Use the search_engine tool to generate four distinct search queries that will help gather a broad range of information related to the user's query. Each query should focus on a different aspect or angle of the topic.",
+	"tools": [
+		{
+			"type": "function",
+			"function": {
+				"name": "search_engine",
+				"description": "A Wikipedia search engine. Generate four distinct search queries to maximize the spread of search results.",
+				"parameters": {
+					"type": "object",
+					"properties": {
+						"queries": {
+							"type": "array",
+							"items": {
+								"type": "string",
+								"description": "A distinct search query focusing on a specific aspect of the topic."
+							},
+							"minItems": 4,
+							"maxItems": 4,
+							"description": "Four distinct search queries to maximize the spread of search results."
+						}
+					},
+					"required": ["queries"]
+				}
+			}
+		}
+	]
+}
diff --git a/config/wiki/wiki.json b/config/wiki/wiki.json
new file mode 100644
index 0000000..b283f2c
--- /dev/null
+++ b/config/wiki/wiki.json
@@ -0,0 +1,4 @@
+{
+	"wiki_url": "http://localhost:1025",
+	"zim_name": "wikipedia_fr_all_maxi_2024-05"
+}
diff --git a/src/app/mod.rs b/src/app.rs
index 3cff678..8376f8c 100644
--- a/src/app/mod.rs
+++ b/src/app.rs
@@ -1,2 +1,3 @@
 pub mod init;
 pub mod llm;
+pub mod modules;
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
             }
diff --git a/src/app/llm.rs b/src/app/llm.rs
index 9c6d222..979b790 100644
--- a/src/app/llm.rs
+++ b/src/app/llm.rs
@@ -15,7 +15,7 @@ pub struct LLM {
 }
 
 impl LLM {
-    pub fn new(config_file: String) -> LLM {
+    pub fn new(config_file: &str) -> LLM {
         let contents = fs::read_to_string(config_file).unwrap();
         serde_json::from_str(&contents).unwrap()
     }
@@ -40,7 +40,7 @@ impl LLM {
                 while let Some(chunk) = res.chunk().await? {
                     let answer: Value = serde_json::from_slice(chunk.as_ref())?;
 
-                    warn(answer.to_string());
+                    //warn(answer.to_string());
                     if answer["done"].as_bool().unwrap_or(false) {
                         break;
                     }
@@ -59,7 +59,7 @@ impl LLM {
 
     // Use tools functionnality of Ollama, only some models supports it:
     // https://ollama.com/search?c=tools
-    pub async fn ask_format(&self, messages: &Vec<Message>) -> Result<serde_json::Value, Box<dyn std::error::Error>> {
+    pub async fn ask_tools(&self, messages: &Vec<Message>) -> Result<serde_json::Value, Box<dyn std::error::Error>> {
         let client = Client::new();
         let response = client
             .post(&self.url)
@@ -72,7 +72,7 @@ impl LLM {
             .send()
             .await?.json::<Value>().await?;
 
-        warn(response.to_string());
+        //warn(response.to_string());
 
         if let Some(tool_calls) = response
             .get("message")
@@ -105,8 +105,8 @@ impl fmt::Display for MessageType {
 
 #[derive(Debug, Serialize, Clone)]
 pub struct Message {
-    role: MessageType,
-    content: String,
+    pub role: MessageType,
+    pub content: String,
 }
 
 impl Message {
@@ -114,6 +114,13 @@ impl Message {
         Message { role, content }
     }
 
+    pub fn default() -> Message {
+        Message {
+            role: MessageType::USER,
+            content: "".to_string(),
+        }
+    }
+
     pub fn save_message(&self, conv_id: String) -> Result<(), Box<dyn std::error::Error>> {
         // Create conv directory if doesn't exist
         create_dir_all("conv")?;
diff --git a/src/app/modules.rs b/src/app/modules.rs
new file mode 100644
index 0000000..622d63c
--- /dev/null
+++ b/src/app/modules.rs
@@ -0,0 +1 @@
+pub mod wikipedia;
diff --git a/src/app/modules/wikipedia.rs b/src/app/modules/wikipedia.rs
new file mode 100644
index 0000000..5864df4
--- /dev/null
+++ b/src/app/modules/wikipedia.rs
@@ -0,0 +1,108 @@
+use crate::app::llm::{Message, MessageType, LLM};
+use percent_encoding::{utf8_percent_encode, NON_ALPHANUMERIC};
+use crate::helper::init::warn;
+use std::fs;
+use select::document::Document;
+use select::predicate::{Name, Class};
+use regex::Regex;
+
+pub async fn ask_wiki(messages: &Vec<Message>) -> Result<String, Box<dyn std::error::Error>> {
+    let wiki_search = LLM::new("config/wiki/wiki-search.json");
+    let wiki_best = LLM::new("config/wiki/wiki-best.json");
+    let wiki_resume = LLM::new("config/wiki/wiki-resume.json");
+
+    let settings: serde_json::Value = serde_json::from_str(&fs::read_to_string("config/wiki/wiki.json").unwrap()).unwrap();
+    let wiki_url: String = settings.get("wiki_url").unwrap().to_string().replace("\"", "");
+    let zim_name: String = settings.get("zim_name").unwrap().to_string().replace("\"", "");
+
+    // Search articles corresponding to user query
+    let user_query: Message = messages.last().unwrap().clone();
+    let articles: Vec<String> = search_articles(user_query.clone(), wiki_search, &wiki_url, &zim_name).await?;
+
+    // Find best article to respond user query
+    let best_article_content = find_get_best_article(articles, &user_query.content, wiki_best, &wiki_url, &zim_name).await?;
+
+    // Resume article and create the response
+    let messages = vec![
+        Message::new(MessageType::SYSTEM, wiki_resume.system_prompt.clone()),
+        Message::new(MessageType::USER, format!("The users query is: {}", user_query.content)),
+        Message::new(MessageType::USER, format!("The search results are: {}", best_article_content)),
+    ];
+    let query_response: String = wiki_resume.ask(&messages).await.unwrap();
+
+    Ok(query_response)
+}
+
+async fn search_articles(user_query: Message, search_llm: LLM, wiki_url: &String, zim_name: &String) -> Result<Vec<String>, Box<dyn std::error::Error>> {
+    // Use LLM to create 4 queries and fetch articles with those 4 queries
+    let messages = vec![
+        Message::new(MessageType::SYSTEM, search_llm.system_prompt.clone()),
+        user_query,
+    ];
+    let result = search_llm.ask_tools(&messages).await?;
+
+    let queries: Vec<String> = result[0]["function"]["arguments"]["queries"].as_array().unwrap().iter().map(|x| x.as_str().unwrap().to_string()).collect();
+
+    // Search articles on wikipedia API
+    let mut articles: Vec<String> = Vec::new();
+    for query in queries.iter() {
+        warn(query.clone());
+
+        // Request kiwix API for articles matching query
+        let encoded_query = utf8_percent_encode(&query, NON_ALPHANUMERIC).to_string(); 
+        let client = reqwest::Client::new();
+        let url = format!("{}/search?books.name={}&pattern={}", wiki_url, zim_name, encoded_query);
+        let body = client.get(url).send().await?.text().await?;
+
+        // Select every article corresponding to the query
+        let document = Document::from(body.as_str());
+
+        // Select articles title from the query
+        let results_div = document.find(Class("results")).next().unwrap();
+        for node in results_div.find(Name("a")) {
+            let article = node.text();
+            articles.push(article.clone());
+        }
+    }
+    Ok(articles)
+}
+
+async fn find_get_best_article(articles: Vec<String>, user_query: &String, best_llm: LLM, wiki_url: &String, zim_name: &String) -> Result<String, Box<dyn std::error::Error>> {
+    // Create a string with all the articles title
+    let mut articles_headings: String = String::new();
+    for article in articles {
+        articles_headings = format!("{}, {}", &articles_headings, article);
+    }
+
+    let messages = vec![
+        Message::new(MessageType::SYSTEM, best_llm.system_prompt.clone()),
+        Message::new(MessageType::USER, format!("The user's query is: {}. Here are the headings:\n{}\n\nPlease select the most relevant heading. Output the heading **only** and nothing else.", user_query, articles_headings))];
+    let best_article = best_llm.ask(&messages).await?;
+
+    // wiki query get article content & parse
+    let client = reqwest::Client::new();
+    let url: String = format!("{}/content/{}/A/{}", wiki_url, zim_name, best_article.replace("*","").replace(" ", "_"));
+    let body = client.get(url).send().await?.text().await?;
+    let content = extract_text_from_tags(&body);
+
+    Ok(content)
+}
+
+fn extract_text_from_tags(html: &str) -> String {
+    // Créer une expression régulière pour trouver le contenu dans les balises <p>, <h1>, <h2>, <h3>
+    let re = Regex::new(r#"<p[^>]*>(.*?)</p>|<h1[^>]*>(.*?)</h1>|<h2[^>]*>(.*?)</h2>|<h3[^>]*>(.*?)</h3>"#).unwrap();
+
+    // Utiliser l'expression régulière pour capturer le contenu des balises <p>, <h1>, <h2>, <h3>
+    let text = re.captures_iter(html)
+        .flat_map(|cap| {
+            // Trouver le premier groupe capturé non vide (parmi cap[1] à cap[4])
+            (1..=4)
+                .filter_map(|i| cap.get(i))
+                .map(|m| m.as_str()) // &str
+                .flat_map(|s| s.split_whitespace())
+                .collect::<Vec<_>>() // Vec<&str>
+        })
+        .collect::<Vec<_>>() // collect words
+        .join(" "); // join with spaces
+    text
+}
diff --git a/src/helper/mod.rs b/src/helper.rs
index 43763f1..43763f1 100644
--- a/src/helper/mod.rs
+++ b/src/helper.rs
diff --git a/src/helper/init.rs b/src/helper/init.rs
index 2e0537d..d004b34 100644
--- a/src/helper/init.rs
+++ b/src/helper/init.rs
@@ -10,5 +10,5 @@ pub fn warn(content: String) {
         .open("log.txt")
         .unwrap();
     let utc: DateTime<Local> = Local::now();
-    writeln!(file, "[{}] {}", utc, content);
+    writeln!(file, "[{}] {}", utc, content).unwrap();
 }
diff --git a/src/lib.rs b/src/lib.rs
deleted file mode 100644
index 52c86c2..0000000
--- a/src/lib.rs
+++ /dev/null
@@ -1,3 +0,0 @@
-pub mod app;
-pub mod helper;
-pub mod ui;
diff --git a/src/ui/mod.rs b/src/ui.rs
index 424376c..424376c 100644
--- a/src/ui/mod.rs
+++ b/src/ui.rs
diff --git a/src/ui/init.rs b/src/ui/init.rs
index ea0882c..afd686a 100644
--- a/src/ui/init.rs
+++ b/src/ui/init.rs
@@ -95,6 +95,10 @@ impl Ui {
         ]);
         let [help_area, messages_area, input_area] = vertical.areas(frame.area());
 
+        let help_horizontal =
+            Layout::horizontal([Constraint::Percentage(75), Constraint::Percentage(25)]);
+        let [help_text_area, conv_id_area] = help_horizontal.areas(help_area);
+
         let (msg, style) = match self.input_field.input_mode {
             InputMode::Normal => (
                 vec![
@@ -103,8 +107,8 @@ impl Ui {
                     " to exit, ".into(),
                     "e".bold(),
                     " to start editing, ".into(),
-                    "r".bold(),
-                    " to resume the conversation.".into(),
+                    "s".bold(),
+                    " to save a resume of the conversation.".into(),
                 ],
                 Style::default(),
             ),
@@ -119,9 +123,13 @@ impl Ui {
                 Style::default(),
             ),
         };
-        let text = Text::from(Line::from(msg)).patch_style(style);
-        let help_message = Paragraph::new(text);
-        frame.render_widget(help_message, help_area);
+        let help_text = Text::from(Line::from(msg)).patch_style(style);
+        let help_message = Paragraph::new(help_text);
+        frame.render_widget(help_message, help_text_area);
+
+        let conv_id = self.app.conv_id.to_string().clone();
+        let conv_id_text = Paragraph::new(format!("Conv id: {conv_id}"));
+        frame.render_widget(conv_id_text, conv_id_area);
 
         // Rendering inputfield
         let input = Paragraph::new(self.input_field.input.as_str())
@@ -187,6 +195,7 @@ impl Ui {
                 max_char_per_line = size;
             }
         }
+
         let messages = Paragraph::new(Text::from(messages))
             .block(Block::bordered().title("Chat with Néo AI"))
             .wrap(Wrap { trim: false })