feat: support RAG (#560)

* feat: support RAG

* support more embeddings models and implement concurrent embedding api

* show the progress of addings paths

* ignore embedding context when saving message

* embedding model max_chunk_size => default_chunk_size

* support pdf and pandoc formats (docx, epub, ipynb)

1 files changed, 425 insertions, 0 deletions

diff --git a/src/rag/mod.rs b/src/rag/mod.rs
new file mode 100644
index 0000000..387d3d9
--- /dev/null
+++ b/src/rag/mod.rs
@@ -0,0 +1,425 @@
+use self::loader::*;
+use self::splitter::*;
+
+use crate::client::*;
+use crate::config::*;
+use crate::utils::*;
+
+mod loader;
+mod splitter;
+
+use anyhow::bail;
+use anyhow::{anyhow, Context, Result};
+use hnsw_rs::prelude::*;
+use indexmap::IndexMap;
+use inquire::{required, validator::Validation, Select, Text};
+use path_absolutize::Absolutize;
+use serde::{Deserialize, Serialize};
+use serde_json::json;
+use std::fmt::Debug;
+use std::{io::BufReader, path::Path};
+use tokio::sync::mpsc;
+
+pub const TEMP_RAG_NAME: &str = "temp";
+pub const CHUNK_OVERLAP: usize = 20;
+pub const SIMILARITY_THRESHOLD: f32 = 0.25;
+
+pub struct Rag {
+    client: Box<dyn Client>,
+    name: String,
+    path: String,
+    model: Model,
+    hnsw: Hnsw<'static, f32, DistCosine>,
+    data: RagData,
+}
+
+impl Debug for Rag {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("Rag")
+            .field("name", &self.name)
+            .field("path", &self.path)
+            .field("model", &self.model)
+            .field("data", &self.data)
+            .finish()
+    }
+}
+
+impl Rag {
+    pub async fn init(
+        config: &GlobalConfig,
+        name: &str,
+        path: &Path,
+        abort_signal: AbortSignal,
+    ) -> Result<Self> {
+        debug!("init rag: {name}");
+        let model = select_embedding_model(config)?;
+        let chunk_size = model.default_chunk_size();
+        let chunk_size = set_chunk_size(chunk_size)?;
+        let data = RagData::new(&model.id(), chunk_size);
+        let mut rag = Self::create(config, name, path, data)?;
+        let paths = add_document_paths()?;
+        debug!("document paths: {paths:?}");
+        let (stop_spinner_tx, set_spinner_message_tx) = run_spinner("Starting").await;
+        tokio::select! {
+            ret = rag.add_paths(&paths, Some(set_spinner_message_tx)) => {
+                let _ = stop_spinner_tx.send(());
+                ret?;
+            }
+            _ = watch_abort_signal(abort_signal) => {
+                let _ = stop_spinner_tx.send(());
+                bail!("Aborted!")
+            },
+        };
+        if !rag.is_temp() {
+            rag.save(path)?;
+            println!("✨ Saved rag to '{}'", path.display());
+        }
+        Ok(rag)
+    }
+
+    pub fn load(config: &GlobalConfig, name: &str, path: &Path) -> Result<Self> {
+        let err = || format!("Failed to load rag '{name}'");
+        let file = std::fs::File::open(path).with_context(err)?;
+        let reader = BufReader::new(file);
+        let data: RagData = bincode::deserialize_from(reader).with_context(err)?;
+        Self::create(config, name, path, data)
+    }
+
+    pub fn create(config: &GlobalConfig, name: &str, path: &Path, data: RagData) -> Result<Self> {
+        let hnsw = data.build_hnsw();
+        let model = retrieve_embedding_model(&config.read(), &data.model)?;
+        let client = init_client(config, Some(model.clone()))?;
+        let rag = Rag {
+            client,
+            name: name.to_string(),
+            path: path.display().to_string(),
+            data,
+            model,
+            hnsw,
+        };
+        Ok(rag)
+    }
+
+    pub fn save(&self, path: &Path) -> Result<()> {
+        ensure_parent_exists(path)?;
+        let mut file = std::fs::File::create(path)?;
+        bincode::serialize_into(&mut file, &self.data)
+            .with_context(|| format!("Failed to save rag '{}'", self.name))?;
+        Ok(())
+    }
+
+    pub fn export(&self) -> Result<String> {
+        let files: Vec<_> = self.data.files.iter().map(|v| &v.path).collect();
+        let data = json!({
+            "path": self.path,
+            "model": self.model.id(),
+            "chunk_size": self.data.chunk_size,
+            "files": files,
+        });
+        let output = serde_yaml::to_string(&data)
+            .with_context(|| format!("Unable to show info about rag '{}'", self.name))?;
+        Ok(output)
+    }
+
+    pub fn name(&self) -> &str {
+        &self.name
+    }
+
+    pub fn is_temp(&self) -> bool {
+        self.name == TEMP_RAG_NAME
+    }
+
+    pub async fn search(
+        &self,
+        text: &str,
+        top_k: usize,
+        abort_signal: AbortSignal,
+    ) -> Result<String> {
+        let (stop_spinner_tx, _) = run_spinner("Embedding").await;
+        let ret = tokio::select! {
+            ret = self.search_impl(text, top_k) => {
+                ret
+            }
+            _ = watch_abort_signal(abort_signal) => {
+                bail!("Aborted!")
+            },
+        };
+        let _ = stop_spinner_tx.send(());
+        let output = ret?.join("\n\n");
+        Ok(output)
+    }
+
+    pub async fn add_paths<T: AsRef<Path>>(
+        &mut self,
+        paths: &[T],
+        progress_tx: Option<mpsc::UnboundedSender<String>>,
+    ) -> Result<()> {
+        // List files
+        let mut file_paths = vec![];
+        progress(&progress_tx, "Listing paths".into());
+        for path in paths {
+            let path = path
+                .as_ref()
+                .absolutize()
+                .with_context(|| anyhow!("Invalid path '{}'", path.as_ref().display()))?;
+            let path_str = path.display().to_string();
+            if self.data.files.iter().any(|v| v.path == path_str) {
+                continue;
+            }
+            let (path_str, suffixes) = parse_glob(&path_str)?;
+            let suffixes = if suffixes.is_empty() {
+                None
+            } else {
+                Some(&suffixes)
+            };
+            list_files(&mut file_paths, Path::new(&path_str), suffixes).await?;
+        }
+
+        // Load files
+        let mut rag_files = vec![];
+        let file_paths_len = file_paths.len();
+        progress(&progress_tx, format!("Loading files [1/{file_paths_len}]"));
+        for path in file_paths {
+            let extension = Path::new(&path)
+                .extension()
+                .map(|v| v.to_string_lossy().to_lowercase())
+                .unwrap_or_default();
+            let separator = autodetect_separator(&extension);
+            let splitter = Splitter::new(self.data.chunk_size, CHUNK_OVERLAP, separator);
+            let documents = load(&path, &extension)
+                .await
+                .with_context(|| format!("Failed to load text at '{path}'"))?;
+            let documents =
+                splitter.split_documents(&documents, &SplitterChunkHeaderOptions::default());
+            rag_files.push(RagFile { path, documents });
+            progress(
+                &progress_tx,
+                format!("Loading files [{}/{file_paths_len}]", rag_files.len()),
+            );
+        }
+
+        if rag_files.is_empty() {
+            return Ok(());
+        }
+
+        // Convert vectors
+        let mut vector_ids = vec![];
+        let mut texts = vec![];
+        for (file_index, file) in rag_files.iter().enumerate() {
+            for (document_index, doc) in file.documents.iter().enumerate() {
+                vector_ids.push(combine_vector_id(file_index, document_index));
+                texts.push(doc.page_content.clone())
+            }
+        }
+
+        let embeddings_data = EmbeddingsData::new(texts, false);
+        let embeddings = self
+            .create_embeddings(embeddings_data, progress_tx.clone())
+            .await?;
+
+        self.data.add(rag_files, vector_ids, embeddings);
+        progress(&progress_tx, "Building vector store".into());
+        self.hnsw = self.data.build_hnsw();
+
+        Ok(())
+    }
+
+    async fn search_impl(&self, text: &str, top_k: usize) -> Result<Vec<String>> {
+        let splitter = Splitter::new(self.data.chunk_size, CHUNK_OVERLAP, &DEFAULT_SEPARATES);
+        let texts = splitter.split_text(text);
+        let embeddings_data = EmbeddingsData::new(texts, true);
+        let embeddings = self.create_embeddings(embeddings_data, None).await?;
+        let output = self
+            .hnsw
+            .parallel_search(&embeddings, top_k, 30)
+            .into_iter()
+            .flat_map(|list| {
+                list.into_iter()
+                    .filter_map(|v| {
+                        if v.distance < SIMILARITY_THRESHOLD {
+                            return None;
+                        }
+                        let (file_index, document_index) = split_vector_id(v.d_id);
+                        let text = self.data.files[file_index].documents[document_index]
+                            .page_content
+                            .clone();
+                        Some(text)
+                    })
+                    .collect::<Vec<_>>()
+            })
+            .collect();
+        Ok(output)
+    }
+
+    async fn create_embeddings(
+        &self,
+        data: EmbeddingsData,
+        progress_tx: Option<mpsc::UnboundedSender<String>>,
+    ) -> Result<EmbeddingsOutput> {
+        let EmbeddingsData { texts, query } = data;
+        let mut output = vec![];
+        let chunks = texts.chunks(self.model.max_concurrent_chunks());
+        let chunks_len = chunks.len();
+        progress(
+            &progress_tx,
+            format!("Creating embeddings [1/{chunks_len}]"),
+        );
+        for (index, texts) in chunks.enumerate() {
+            let chunk_data = EmbeddingsData {
+                texts: texts.to_vec(),
+                query,
+            };
+            let chunk_output = self
+                .client
+                .embeddings(chunk_data)
+                .await
+                .context("Failed to create embedding")?;
+            output.extend(chunk_output);
+            progress(
+                &progress_tx,
+                format!("Creating embeddings [{}/{chunks_len}]", index + 1),
+            );
+        }
+        Ok(output)
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct RagData {
+    pub model: String,
+    pub chunk_size: usize,
+    pub files: Vec<RagFile>,
+    pub vectors: IndexMap<VectorID, Vec<f32>>,
+}
+
+impl RagData {
+    pub fn new(model: &str, chunk_size: usize) -> Self {
+        Self {
+            model: model.to_string(),
+            chunk_size,
+            files: Default::default(),
+            vectors: Default::default(),
+        }
+    }
+
+    pub fn add(
+        &mut self,
+        files: Vec<RagFile>,
+        vector_ids: Vec<VectorID>,
+        embeddings: EmbeddingsOutput,
+    ) {
+        self.files.extend(files);
+        self.vectors.extend(vector_ids.into_iter().zip(embeddings));
+    }
+
+    pub fn build_hnsw(&self) -> Hnsw<'static, f32, DistCosine> {
+        let hnsw = Hnsw::new(32, self.vectors.len(), 16, 200, DistCosine {});
+        let list: Vec<_> = self.vectors.iter().map(|(k, v)| (v, *k)).collect();
+        hnsw.parallel_insert(&list);
+        hnsw
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct RagFile {
+    path: String,
+    documents: Vec<RagDocument>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct RagDocument {
+    pub page_content: String,
+    pub metadata: RagMetadata,
+}
+
+impl RagDocument {
+    pub fn new<S: Into<String>>(page_content: S) -> Self {
+        RagDocument {
+            page_content: page_content.into(),
+            metadata: IndexMap::new(),
+        }
+    }
+
+    #[allow(unused)]
+    pub fn with_metadata(mut self, metadata: RagMetadata) -> Self {
+        self.metadata = metadata;
+        self
+    }
+}
+
+impl Default for RagDocument {
+    fn default() -> Self {
+        RagDocument {
+            page_content: "".to_string(),
+            metadata: IndexMap::new(),
+        }
+    }
+}
+
+pub type RagMetadata = IndexMap<String, String>;
+
+pub type VectorID = usize;
+
+pub fn combine_vector_id(file_index: usize, document_index: usize) -> VectorID {
+    file_index << (usize::BITS / 2) | document_index
+}
+
+pub fn split_vector_id(value: VectorID) -> (usize, usize) {
+    let low_mask = (1 << (usize::BITS / 2)) - 1;
+    let low = value & low_mask;
+    let high = value >> (usize::BITS / 2);
+    (high, low)
+}
+
+fn retrieve_embedding_model(config: &Config, model_id: &str) -> Result<Model> {
+    let models = list_embedding_models(config);
+    let model =
+        Model::find(&models, model_id).ok_or_else(|| anyhow!("No embedding model '{model_id}'"))?;
+    Ok(model)
+}
+
+fn select_embedding_model(config: &GlobalConfig) -> Result<Model> {
+    let config = config.read();
+    let model = match config.embedding_model.clone() {
+        Some(model_id) => retrieve_embedding_model(&config, &model_id)?,
+        None => {
+            let models = list_embedding_models(&config);
+            if models.is_empty() {
+                bail!("No embedding model");
+            }
+            let model_ids: Vec<_> = models.iter().map(|v| v.id()).collect();
+            let model_id = Select::new("Select embedding model:", model_ids).prompt()?;
+            retrieve_embedding_model(&config, &model_id)?
+        }
+    };
+    Ok(model)
+}
+
+fn set_chunk_size(chunk_size: usize) -> Result<usize> {
+    let value = Text::new("Set chunk size:")
+        .with_default(&chunk_size.to_string())
+        .with_validator(move |text: &str| {
+            let out = match text.parse::<usize>() {
+                Ok(_) => Validation::Valid,
+                Err(_) => Validation::Invalid("Must be a integer".into()),
+            };
+            Ok(out)
+        })
+        .prompt()?;
+    value.parse().map_err(|_| anyhow!("Invalid chunk_size"))
+}
+
+fn add_document_paths() -> Result<Vec<String>> {
+    let text = Text::new("Add document paths:")
+        .with_validator(required!("This field is required"))
+        .with_help_message("e.g. file1;dir2/;dir3/**/*.md")
+        .prompt()?;
+    let paths = text.split(';').map(|v| v.to_string()).collect();
+    Ok(paths)
+}
+
+fn progress(spinner_message_tx: &Option<mpsc::UnboundedSender<String>>, message: String) {
+    if let Some(tx) = spinner_message_tx {
+        let _ = tx.send(message);
+    }
+}