summaryrefslogtreecommitdiffstats
path: root/src/rag/mod.rs
diff options
context:
space:
mode:
Diffstat (limited to 'src/rag/mod.rs')
-rw-r--r--src/rag/mod.rs126
1 files changed, 71 insertions, 55 deletions
diff --git a/src/rag/mod.rs b/src/rag/mod.rs
index e84ece3..4cee391 100644
--- a/src/rag/mod.rs
+++ b/src/rag/mod.rs
@@ -1,4 +1,3 @@
-use self::bm25::*;
use self::loader::*;
use self::splitter::*;
@@ -6,12 +5,12 @@ use crate::client::*;
use crate::config::*;
use crate::utils::*;
-mod bm25;
mod loader;
mod serde_vectors;
mod splitter;
use anyhow::{anyhow, bail, Context, Result};
+use bm25::{LanguageMode, SearchEngine, SearchEngineBuilder};
use hnsw_rs::prelude::*;
use indexmap::{IndexMap, IndexSet};
use inquire::{required, validator::Validation, Confirm, Select, Text};
@@ -19,7 +18,7 @@ use parking_lot::RwLock;
use path_absolutize::Absolutize;
use serde::{Deserialize, Serialize};
use serde_json::json;
-use std::{collections::HashMap, env, fmt::Debug, fs, path::Path, time::Duration};
+use std::{collections::HashMap, env, fmt::Debug, fs, hash::Hash, path::Path, time::Duration};
use tokio::time::sleep;
pub struct Rag {
@@ -28,7 +27,7 @@ pub struct Rag {
path: String,
embedding_model: Model,
hnsw: Hnsw<'static, f32, DistCosine>,
- bm25: BM25<DocumentId>,
+ bm25: SearchEngine<DocumentId>,
data: RagData,
last_sources: RwLock<Option<String>>,
}
@@ -52,7 +51,7 @@ impl Clone for Rag {
path: self.path.clone(),
embedding_model: self.embedding_model.clone(),
hnsw: self.data.build_hnsw(),
- bm25: self.bm25.clone(),
+ bm25: self.data.build_bm25(),
data: self.data.clone(),
last_sources: RwLock::new(None),
}
@@ -230,7 +229,7 @@ impl Rag {
let sources: IndexSet<_> = ids
.iter()
.filter_map(|id| {
- let (file_index, _) = split_document_id(*id);
+ let (file_index, _) = id.split();
let file = self.data.files.get(&file_index)?;
Some(file.path.clone())
})
@@ -394,14 +393,7 @@ impl Rag {
&separator,
);
- let metadata = metadata
- .iter()
- .map(|(k, v)| format!("{k}: {v}\n"))
- .collect::<Vec<String>>()
- .join("");
- let split_options = SplitterChunkHeaderOptions::default().with_chunk_header(&format!(
- "<document_metadata>\npath: {path}\n{metadata}</document_metadata>\n\n"
- ));
+ let split_options = SplitterChunkHeaderOptions::default();
let document = RagDocument::new(contents);
let split_documents = splitter.split_documents(&[document], &split_options);
rag_files.push(RagFile {
@@ -420,7 +412,7 @@ impl Rag {
let mut texts = vec![];
for file in rag_files.into_iter() {
for (document_index, document) in file.documents.iter().enumerate() {
- document_ids.push(combine_document_id(next_file_id, document_index));
+ document_ids.push(DocumentId::new(next_file_id, document_index));
texts.push(document.page_content.clone())
}
files.push((next_file_id, file));
@@ -456,17 +448,21 @@ impl Rag {
min_score_keyword_search: f32,
rerank_model: Option<&str>,
) -> Result<Vec<(DocumentId, String)>> {
- let (vector_search_result, text_search_result) = tokio::join!(
+ let (vector_search_results, keyword_search_results) = tokio::join!(
self.vector_search(query, top_k, min_score_vector_search),
self.keyword_search(query, top_k, min_score_keyword_search)
);
- let vector_search_ids = vector_search_result?;
- let keyword_search_ids = text_search_result?;
- debug!(
- "vector_search_ids: {:?}, keyword_search_ids: {:?}",
- pretty_document_ids(&vector_search_ids),
- pretty_document_ids(&keyword_search_ids)
- );
+
+ let vector_search_results = vector_search_results?;
+ debug!("vector_search_results: {vector_search_results:?}",);
+ let vector_search_ids: Vec<DocumentId> =
+ vector_search_results.into_iter().map(|(v, _)| v).collect();
+
+ let keyword_search_results = keyword_search_results?;
+ debug!("keyword_search_results: {keyword_search_results:?}",);
+ let keyword_search_ids: Vec<DocumentId> =
+ keyword_search_results.into_iter().map(|(v, _)| v).collect();
+
let ids = match rerank_model {
Some(model_id) => {
let model = Model::retrieve_reranker(&self.config.read(), model_id)?;
@@ -490,7 +486,7 @@ impl Rag {
.take(top_k)
.filter_map(|item| documents_ids.get(item.index).cloned())
.collect();
- debug!("rerank_ids: {:?}", pretty_document_ids(&ids));
+ debug!("rerank_ids: {ids:?}");
ids
}
None => {
@@ -499,7 +495,7 @@ impl Rag {
vec![1.0, 1.0],
top_k,
);
- debug!("rrf_ids: {:?}", pretty_document_ids(&ids));
+ debug!("rrf_ids: {ids:?}");
ids
}
};
@@ -518,7 +514,7 @@ impl Rag {
query: &str,
top_k: usize,
min_score: f32,
- ) -> Result<Vec<DocumentId>> {
+ ) -> Result<Vec<(DocumentId, f32)>> {
let splitter = RecursiveCharacterTextSplitter::new(
self.data.chunk_size,
self.data.chunk_overlap,
@@ -534,10 +530,12 @@ impl Rag {
.flat_map(|list| {
list.into_iter()
.filter_map(|v| {
- if v.distance < min_score {
- return None;
+ let score = 1.0 - v.distance;
+ if score > min_score {
+ Some((DocumentId(v.d_id), score))
+ } else {
+ None
}
- Some(v.d_id)
})
.collect::<Vec<_>>()
})
@@ -550,8 +548,19 @@ impl Rag {
query: &str,
top_k: usize,
min_score: f32,
- ) -> Result<Vec<DocumentId>> {
- let output = self.bm25.search(query, top_k, Some(min_score as f64));
+ ) -> Result<Vec<(DocumentId, f32)>> {
+ let results = self.bm25.search(query, top_k);
+ let output: Vec<(DocumentId, f32)> = results
+ .into_iter()
+ .filter_map(|v| {
+ let score = v.score;
+ if score > min_score {
+ Some((v.document.id, score))
+ } else {
+ None
+ }
+ })
+ .collect();
Ok(output)
}
@@ -670,7 +679,7 @@ impl RagData {
}
pub fn get(&self, id: DocumentId) -> Option<&RagDocument> {
- let (file_index, document_index) = split_document_id(id);
+ let (file_index, document_index) = id.split();
let file = self.files.get(&file_index)?;
let document = file.documents.get(document_index)?;
Some(document)
@@ -680,7 +689,7 @@ impl RagData {
for file_id in file_ids {
if let Some(file) = self.files.swap_remove(&file_id) {
for (document_index, _) in file.documents.iter().enumerate() {
- let document_id = combine_document_id(file_id, document_index);
+ let document_id = DocumentId::new(file_id, document_index);
self.vectors.swap_remove(&document_id);
}
}
@@ -702,20 +711,23 @@ impl RagData {
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();
+ let list: Vec<_> = self.vectors.iter().map(|(k, v)| (v, k.0)).collect();
hnsw.parallel_insert(&list);
hnsw
}
- pub fn build_bm25(&self) -> BM25<DocumentId> {
- let mut corpus = vec![];
+ pub fn build_bm25(&self) -> SearchEngine<DocumentId> {
+ let mut documents = vec![];
for (file_index, file) in self.files.iter() {
for (document_index, document) in file.documents.iter().enumerate() {
- let id = combine_document_id(*file_index, document_index);
- corpus.push((id, document.page_content.clone()));
+ let id = DocumentId::new(*file_index, document_index);
+ documents.push(bm25::Document::new(id, &document.page_content))
}
}
- BM25::new(corpus, BM25Options::default())
+ SearchEngineBuilder::<DocumentId>::with_documents(LanguageMode::Detect, documents)
+ .k1(1.5)
+ .b(0.75)
+ .build()
}
}
@@ -753,26 +765,30 @@ impl Default for RagDocument {
pub type RagMetadata = IndexMap<String, String>;
pub type FileId = usize;
-pub type DocumentId = usize;
-pub fn combine_document_id(file_index: usize, document_index: usize) -> DocumentId {
- file_index << (usize::BITS / 2) | document_index
-}
+#[derive(Clone, Copy, Hash, Eq, PartialEq, Ord, PartialOrd)]
+pub struct DocumentId(usize);
-pub fn split_document_id(value: DocumentId) -> (usize, usize) {
- let low_mask = (1 << (usize::BITS / 2)) - 1;
- let low = value & low_mask;
- let high = value >> (usize::BITS / 2);
- (high, low)
+impl Debug for DocumentId {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ let (file_index, document_index) = self.split();
+ f.write_fmt(format_args!("{file_index}-{document_index}"))
+ }
}
-fn pretty_document_ids(ids: &[DocumentId]) -> Vec<String> {
- ids.iter()
- .map(|v| {
- let (h, l) = split_document_id(*v);
- format!("{h}-{l}")
- })
- .collect()
+impl DocumentId {
+ pub fn new(file_index: usize, document_index: usize) -> Self {
+ let value = file_index << (usize::BITS / 2) | document_index;
+ Self(value)
+ }
+
+ pub fn split(self) -> (usize, usize) {
+ let value = self.0;
+ let low_mask = (1 << (usize::BITS / 2)) - 1;
+ let low = value & low_mask;
+ let high = value >> (usize::BITS / 2);
+ (high, low)
+ }
}
fn select_embedding_model(models: &[&Model]) -> Result<String> {
generated by cgit v1.2.3 (git 2.39.1) at 2026-10-03 00:46:16 +0000