RAG implementation

crates/server/src/main.rs

@@ -1,10 +1,14 @@
-use std::sync::{Arc, LazyLock};
+use std::{
+    path::PathBuf,
+    sync::{Arc, LazyLock},
+};
 
 use miniserve::{http::StatusCode, Content, Request, Response};
 use serde::{Deserialize, Serialize};
 use tokio::{
-    join,
+    fs, join,
     sync::{mpsc, oneshot},
+    task::JoinSet,
 };
 
 async fn index(_req: Request) -> Response {
@@ -17,19 +21,36 @@ struct Messages {
     messages: Vec<String>,
 }
 
-async fn query_chat(messages: &Arc<Vec<String>>) -> Vec<String> {
-    type Payload = (Arc<Vec<String>>, oneshot::Sender<Vec<String>>);
-    static SENDER: LazyLock<mpsc::Sender<Payload>> = LazyLock::new(|| {
-        let (tx, mut rx) = mpsc::channel::<Payload>(1024);
-        tokio::spawn(async move {
-            let mut chatbot = chatbot::Chatbot::new(vec![":-)".into(), "^^".into()]);
-            while let Some((messages, responder)) = rx.recv().await {
-                let response = chatbot.query_chat(&messages).await;
-                responder.send(response).unwrap();
-            }
-        });
-        tx
+async fn load_docs(paths: Vec<PathBuf>) -> Vec<String> {
+    let mut doc_futs = paths
+        .into_iter()
+        .map(fs::read_to_string)
+        .collect::<JoinSet<_>>();
+    let mut docs = Vec::new();
+    while let Some(result) = doc_futs.join_next().await {
+        docs.push(result.unwrap().unwrap());
+    }
+    docs
+}
+
+type Payload = (Arc<Vec<String>>, oneshot::Sender<Vec<String>>);
+
+fn chatbot_thread() -> mpsc::Sender<Payload> {
+    let (tx, mut rx) = mpsc::channel::<Payload>(1024);
+    tokio::spawn(async move {
+        let mut chatbot = chatbot::Chatbot::new(vec![":-)".into(), "^^".into()]);
+        while let Some((messages, responder)) = rx.recv().await {
+            let doc_paths = chatbot.retrieval_documents(&messages);
+            let docs = load_docs(doc_paths).await;
+            let response = chatbot.query_chat(&messages, &docs).await;
+            responder.send(response).unwrap();
+        }
     });
+    tx
+}
+
+async fn query_chat(messages: &Arc<Vec<String>>) -> Vec<String> {
+    static SENDER: LazyLock<mpsc::Sender<Payload>> = LazyLock::new(chatbot_thread);
 
     let (tx, rx) = oneshot::channel();
     SENDER.send((Arc::clone(messages), tx)).await.unwrap();