mcp-memory/server/src/search.rs

use crate::models::{Adr, Entity, Snippet, Task};
use std::sync::Mutex;
use tantivy::schema::*;
use tantivy::{Index, IndexReader, IndexWriter, ReloadPolicy, doc};

pub struct MemoryIndex {
    index: Index,
    reader: IndexReader,
    writer: Mutex<IndexWriter>,

    // Schema fields
    pub id_field: Field,
    pub title_field: Field,
    pub body_field: Field,
    pub type_field: Field,
    pub namespace_field: Field,
}

impl MemoryIndex {
    pub fn new(store_dir: &std::path::Path) -> tantivy::Result<Self> {
        let mut schema_builder = Schema::builder();
        let id_field = schema_builder.add_text_field("id", STRING | STORED);
        let title_field = schema_builder.add_text_field("title", TEXT | STORED);
        let body_field = schema_builder.add_text_field("body", TEXT | STORED);
        let type_field = schema_builder.add_text_field("type", STRING | STORED);
        let namespace_field = schema_builder.add_text_field("namespace", STRING | STORED);
        let schema = schema_builder.build();

        let index_dir = store_dir.join("tantivy_index");
        std::fs::create_dir_all(&index_dir).unwrap();
        let index = Index::open_in_dir(&index_dir).unwrap_or_else(|_| Index::create_in_dir(&index_dir, schema.clone()).unwrap());

        let writer = index.writer(50_000_000)?;
        let reader = index
            .reader_builder()
            .reload_policy(ReloadPolicy::OnCommitWithDelay)
            .try_into()?;

        Ok(Self {
            index,
            reader,
            writer: Mutex::new(writer),
            id_field,
            title_field,
            body_field,
            type_field,
            namespace_field,
        })
    }

    pub fn index_entity(&self, e: &Entity) -> tantivy::Result<()> {
        let mut writer = self.writer.lock().unwrap();
        writer.add_document(doc!(
            self.id_field => e.name.clone(),
            self.title_field => e.name.clone(),
            self.body_field => e.observations.join(" "),
            self.type_field => "entity",
            self.namespace_field => e.namespace.clone()
        ))?;
        writer.commit()?;
        Ok(())
    }

    pub fn index_task(&self, t: &Task) -> tantivy::Result<()> {
        let mut writer = self.writer.lock().unwrap();
        writer.add_document(doc!(
            self.id_field => t.id.clone(),
            self.title_field => t.title.clone(),
            self.body_field => t.description.clone(),
            self.type_field => "task",
            self.namespace_field => "global"
        ))?;
        writer.commit()?;
        Ok(())
    }

    pub fn search(
        &self,
        query: &str,
        namespace: Option<&str>,
    ) -> tantivy::Result<Vec<(String, String)>> {
        let searcher = self.reader.searcher();
        let query_parser = tantivy::query::QueryParser::for_index(
            &self.index,
            vec![self.title_field, self.body_field],
        );
        let q = query_parser.parse_query(query)?;

        let top_docs = searcher.search(
            &q,
            &tantivy::collector::TopDocs::with_limit(50).order_by_score(),
        )?;
        let mut results = Vec::new();
        for (_score, doc_address) in top_docs {
            let retrieved_doc = searcher.doc::<tantivy::TantivyDocument>(doc_address)?;
            let id = retrieved_doc
                .get_first(self.id_field)
                .and_then(|v| v.as_str())
                .unwrap_or("")
                .to_string();
            let doc_type = retrieved_doc
                .get_first(self.type_field)
                .and_then(|v| v.as_str())
                .unwrap_or("")
                .to_string();
            let doc_ns = retrieved_doc
                .get_first(self.namespace_field)
                .and_then(|v| v.as_str())
                .unwrap_or("");
            if let Some(ns) = namespace
                && doc_ns != ns && doc_ns != "global" {
                    continue;
                }
            results.push((id, doc_type));
        }
        Ok(results)
    }

    pub fn index_snippet(&self, s: &Snippet) -> tantivy::Result<()> {
        let mut writer = self.writer.lock().unwrap();
        writer.add_document(doc!(
            self.id_field => s.name.clone(),
            self.title_field => s.name.clone(),
            self.body_field => format!("{} {}", s.language, s.description),
            self.type_field => "snippet",
            self.namespace_field => "global"
        ))?;
        writer.commit()?;
        Ok(())
    }

    pub fn index_adr(&self, a: &Adr) -> tantivy::Result<()> {
        let mut writer = self.writer.lock().unwrap();
        writer.add_document(doc!(
            self.id_field => a.id.clone(),
            self.title_field => a.title.clone(),
            self.body_field => format!("{} {} {}", a.context, a.decision, a.consequence),
            self.type_field => "adr",
            self.namespace_field => "global"
        ))?;
        writer.commit()?;
        Ok(())
    }
}