Create indexer.rs

src/indexer.rs

@@ -0,0 +1,236 @@
+use crate::parser::ParsedDocument;
+use anyhow::{Context, Result};
+use chrono::Utc;
+use serde::{Deserialize, Serialize};
+use std::fs;
+use std::path::{Path, PathBuf};
+use std::sync::Arc;
+use tantivy::collector::TopDocs;
+use tantivy::query::{BooleanQuery, Occur, Query, QueryParser, TermQuery};
+use tantivy::schema::*;
+use tantivy::snippet::SnippetGenerator;
+use tantivy::{doc, Index, IndexReader, IndexWriter, ReloadPolicy, TantivyDocument, Term};
+use tokio::sync::RwLock;
+
+#[derive(Clone)]
+pub struct IndexSchemaFields {
+    pub id: Field,
+    pub collection: Field,
+    pub path: Field,
+    pub title: Field,
+    pub content: Field,
+    pub extension: Field,
+    pub file_size: Field,
+    pub indexed_at: Field,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct SearchHit {
+    pub score: f32,
+    pub collection: String,
+    pub path: String,
+    pub title: String,
+    pub snippet: String,
+    pub extension: String,
+    pub file_size: u64,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct CollectionStats {
+    pub collection_name: String,
+    pub doc_count: u64,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct IndexStats {
+    pub total_documents: u64,
+    pub num_segments: usize,
+    pub index_path: String,
+    pub collections: Vec<CollectionStats>,
+}
+
+pub struct TantivyEngine {
+    index: Index,
+    reader: IndexReader,
+    writer: Arc<RwLock<IndexWriter>>,
+    fields: IndexSchemaFields,
+    base_path: PathBuf,
+}
+
+impl TantivyEngine {
+    pub fn new<P: AsRef<Path>>(base_path: P) -> Result<Self> {
+        let path_buf = base_path.as_ref().to_path_buf();
+        fs::create_dir_all(&path_buf)?;
+
+        let mut schema_builder = Schema::builder();
+        let id = schema_builder.add_text_field("id", STRING | STORED);
+        let collection = schema_builder.add_text_field("collection", STRING | STORED | FAST);
+        let path = schema_builder.add_text_field("path", STRING | STORED);
+        let title = schema_builder.add_text_field("title", TEXT | STORED);
+        let content = schema_builder.add_text_field("content", TEXT | STORED);
+        let extension = schema_builder.add_text_field("extension", STRING | STORED);
+        let file_size = schema_builder.add_u64_field("file_size", STORED | FAST);
+        let indexed_at = schema_builder.add_i64_field("indexed_at", STORED | FAST);
+
+        let schema = schema_builder.build();
+
+        let index = Index::open_or_create(
+            tantivy::directory::MmapDirectory::open(&path_buf)?,
+            schema.clone(),
+        )?;
+
+        // 100MB memory budget for the writer buffer
+        let writer = index.writer(100 * 1024 * 1024)?;
+        let reader = index
+            .reader_builder()
+            .reload_policy(ReloadPolicy::OnCommitWithDelay)
+            .try_into()?;
+
+        let fields = IndexSchemaFields {
+            id,
+            collection,
+            path,
+            title,
+            content,
+            extension,
+            file_size,
+            indexed_at,
+        };
+
+        Ok(Self {
+            index,
+            reader,
+            writer: Arc::new(RwLock::new(writer)),
+            fields,
+            base_path: path_buf,
+        })
+    }
+
+    pub async fn add_documents(&self, collection_name: &str, docs: Vec<ParsedDocument>) -> Result<usize> {
+        let count = docs.len();
+        let writer = self.writer.write().await;
+
+        for doc in docs {
+            let mut tantivy_doc = TantivyDocument::default();
+            let doc_id = uuid::Uuid::new_v4().to_string();
+
+            tantivy_doc.add_text(self.fields.id, &doc_id);
+            tantivy_doc.add_text(self.fields.collection, collection_name);
+            tantivy_doc.add_text(self.fields.path, &doc.path);
+            tantivy_doc.add_text(self.fields.title, &doc.title);
+            tantivy_doc.add_text(self.fields.content, &doc.content);
+            tantivy_doc.add_text(self.fields.extension, &doc.extension);
+            tantivy_doc.add_u64(self.fields.file_size, doc.size_bytes);
+            tantivy_doc.add_i64(self.fields.indexed_at, Utc::now().timestamp());
+
+            writer.add_document(tantivy_doc)?;
+        }
+
+        // Commit and reload index reader
+        let mut writer_guard = writer;
+        writer_guard.commit()?;
+        drop(writer_guard);
+
+        self.reader.reload()?;
+        Ok(count)
+    }
+
+    pub fn search(
+        &self,
+        query_str: &str,
+        collection_filter: Option<&str>,
+        limit: usize,
+    ) -> Result<Vec<SearchHit>> {
+        let searcher = self.reader.searcher();
+        let query_parser = QueryParser::for_index(
+            &self.index,
+            vec![self.fields.title, self.fields.content],
+        );
+
+        let parsed_query = query_parser
+            .parse_query(query_str)
+            .context("Query parse error")?;
+
+        let final_query: Box<dyn Query> = if let Some(coll) = collection_filter {
+            let coll_term = Term::from_field_text(self.fields.collection, coll);
+            let coll_query = TermQuery::new(coll_term, IndexRecordOption::Basic);
+
+            Box::new(BooleanQuery::new(vec![
+                (Occur::Must, parsed_query),
+                (Occur::Must, Box::new(coll_query)),
+            ]))
+        } else {
+            parsed_query
+        };
+
+        let top_docs = searcher.search(&final_query, &TopDocs::with_limit(limit))?;
+        let mut snippet_generator =
+            SnippetGenerator::create(&searcher, &*final_query, self.fields.content)?;
+        snippet_generator.set_max_num_chars(250);
+
+        let mut hits = Vec::with_capacity(top_docs.len());
+
+        for (score, doc_address) in top_docs {
+            let retrieved_doc: TantivyDocument = searcher.doc(doc_address)?;
+
+            let coll_val = retrieved_doc
+                .get_first(self.fields.collection)
+                .and_then(|v| v.as_str())
+                .unwrap_or("default")
+                .to_string();
+
+            let path_val = retrieved_doc
+                .get_first(self.fields.path)
+                .and_then(|v| v.as_str())
+                .unwrap_or("")
+                .to_string();
+
+            let title_val = retrieved_doc
+                .get_first(self.fields.title)
+                .and_then(|v| v.as_str())
+                .unwrap_or("")
+                .to_string();
+
+            let ext_val = retrieved_doc
+                .get_first(self.fields.extension)
+                .and_then(|v| v.as_str())
+                .unwrap_or("")
+                .to_string();
+
+            let size_val = retrieved_doc
+                .get_first(self.fields.file_size)
+                .and_then(|v| v.as_u64())
+                .unwrap_or(0);
+
+            let snippet = snippet_generator.snippet_from_doc(&retrieved_doc).to_html();
+
+            hits.push(SearchHit {
+                score,
+                collection: coll_val,
+                path: path_val,
+                title: title_val,
+                snippet,
+                extension: ext_val,
+                file_size: size_val,
+            });
+        }
+
+        Ok(hits)
+    }
+
+    pub fn get_stats(&self) -> Result<IndexStats> {
+        let searcher = self.reader.searcher();
+        let total_docs = searcher.num_docs();
+        let num_segments = searcher.segment_readers().len();
+
+        Ok(IndexStats {
+            total_documents: total_docs,
+            num_segments,
+            index_path: self.base_path.to_string_lossy().to_string(),
+            collections: vec![CollectionStats {
+                collection_name: "global".to_string(),
+                doc_count: total_docs,
+            }],
+        })
+    }
+}