Skip to main content

semantic_memory_mcp/
bridge.rs

1//! Bridge between the MCP server and semantic-memory MemoryStore.
2//!
3//! semantic-memory's MemoryStore is opened synchronously and search methods
4//! are async (tokio). This bridge wraps the store and uses the current
5//! tokio runtime handle for async calls (no separate runtime).
6
7#[cfg(feature = "candle-embedder")]
8use semantic_memory::embedder::CandleEmbedder;
9use semantic_memory::embedder::{Embedder, MockEmbedder, OllamaEmbedder};
10use semantic_memory::{EmbeddingConfig, MemoryConfig, MemoryStore, SearchConfig};
11use std::path::PathBuf;
12use tokio::runtime::Handle;
13
14/// Which embedding backend to use.
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub enum EmbedderBackend {
17    /// In-process Candle embedder (pure-Rust, CPU-only, no Ollama required).
18    /// Downloads the model from HuggingFace Hub on first use.
19    Candle,
20    /// External Ollama server (requires `ollama serve` running).
21    Ollama,
22    /// Mock embedder for testing (deterministic hash-based, no real embeddings).
23    Mock,
24}
25
26impl Default for EmbedderBackend {
27    fn default() -> Self {
28        // Candle is the default when the feature is enabled.
29        // Ollama is the default when it's not (backward compat).
30        #[cfg(feature = "candle-embedder")]
31        {
32            EmbedderBackend::Candle
33        }
34        #[cfg(not(feature = "candle-embedder"))]
35        {
36            EmbedderBackend::Ollama
37        }
38    }
39}
40
41impl std::str::FromStr for EmbedderBackend {
42    type Err = String;
43
44    fn from_str(s: &str) -> Result<Self, Self::Err> {
45        match s.to_lowercase().as_str() {
46            "candle" => Ok(EmbedderBackend::Candle),
47            "ollama" => Ok(EmbedderBackend::Ollama),
48            "mock" => Ok(EmbedderBackend::Mock),
49            other => Err(format!(
50                "unknown embedder '{other}', expected: candle, ollama, or mock"
51            )),
52        }
53    }
54}
55
56#[derive(Clone)]
57pub struct MemoryBridge {
58    pub store: MemoryStore,
59}
60
61pub struct BridgeConfig {
62    /// MCP-005: renamed from db_path to memory_dir — this is a directory,
63    /// not a SQLite file path. semantic-memory creates base_dir/memory.db
64    /// inside it.
65    pub memory_dir: PathBuf,
66    pub embedder_backend: EmbedderBackend,
67    /// Ollama URL — only used when backend is Ollama.
68    pub embedding_url: String,
69    /// Embedding model name — used by both Candle (as HF model ID) and Ollama.
70    pub embedding_model: String,
71    pub embedding_dims: usize,
72    /// Enable TurboQuant compressed vector candidate backend.
73    pub turbo_quant_enabled: bool,
74    /// TurboQuant polar angle bits (default: 8).
75    pub turbo_quant_bits: Option<u8>,
76    /// TurboQuant QJL projection count (default: 16).
77    pub turbo_quant_projections: Option<usize>,
78}
79
80impl BridgeConfig {
81    pub fn from_args(
82        memory_dir: &str,
83        embedder_backend: EmbedderBackend,
84        embedding_url: Option<&str>,
85        embedding_model: Option<&str>,
86        embedding_dims: Option<usize>,
87        turbo_quant_enabled: bool,
88        turbo_quant_bits: Option<u8>,
89        turbo_quant_projections: Option<usize>,
90    ) -> Self {
91        Self {
92            memory_dir: PathBuf::from(memory_dir),
93            embedder_backend,
94            embedding_url: embedding_url
95                .unwrap_or("http://localhost:11434")
96                .to_string(),
97            embedding_model: embedding_model.unwrap_or("nomic-embed-text").to_string(),
98            embedding_dims: embedding_dims.unwrap_or(768),
99            turbo_quant_enabled,
100            turbo_quant_bits,
101            turbo_quant_projections,
102        }
103    }
104}
105
106impl MemoryBridge {
107    /// Open the memory store with the given config.
108    /// Store opening is synchronous — no runtime needed here.
109    pub fn open(config: BridgeConfig) -> anyhow::Result<Self> {
110        let embedding_config = EmbeddingConfig {
111            ollama_url: config.embedding_url,
112            model: config.embedding_model,
113            dimensions: config.embedding_dims,
114            batch_size: 32,
115            timeout_secs: 60,
116        };
117
118        let embedder: Box<dyn Embedder> = match config.embedder_backend {
119            EmbedderBackend::Mock => Box::new(MockEmbedder::new(config.embedding_dims)),
120            EmbedderBackend::Ollama => Box::new(OllamaEmbedder::try_new(&embedding_config)?),
121            EmbedderBackend::Candle => {
122                #[cfg(feature = "candle-embedder")]
123                {
124                    // For Candle, the model field is the HuggingFace model ID.
125                    // Map common Ollama model names to HF model IDs.
126                    let hf_model_id = match embedding_config.model.as_str() {
127                        "nomic-embed-text" | "nomic-embed-text-v1.5" => {
128                            "nomic-ai/nomic-embed-text-v1.5"
129                        }
130                        other => other,
131                    };
132                    let candle_config = EmbeddingConfig {
133                        model: hf_model_id.to_string(),
134                        ..embedding_config.clone()
135                    };
136                    Box::new(CandleEmbedder::try_new(&candle_config)?)
137                }
138                #[cfg(not(feature = "candle-embedder"))]
139                {
140                    anyhow::bail!(
141                        "candle embedder requested but the 'candle-embedder' feature \
142                         is not enabled. Rebuild with --features candle-embedder \
143                         or use --embedder ollama"
144                    )
145                }
146            }
147        };
148
149        let mut search_config = SearchConfig::default();
150
151        // TurboQuant compressed vector candidate backend
152        #[cfg(feature = "full")]
153        {
154            if config.turbo_quant_enabled {
155                use semantic_memory::DerivedVectorBackendPolicy;
156                search_config.derived_vector_backend =
157                    DerivedVectorBackendPolicy::TurboQuantCandidateOnly;
158                if let Some(bits) = config.turbo_quant_bits {
159                    search_config.turbo_quant_bits = bits;
160                }
161                if let Some(projs) = config.turbo_quant_projections {
162                    search_config.turbo_quant_projections = projs;
163                }
164            }
165        }
166
167        let mem_config = MemoryConfig {
168            base_dir: config.memory_dir,
169            embedding: embedding_config,
170            search: search_config,
171            ..Default::default()
172        };
173
174        let store = MemoryStore::open_with_embedder(mem_config, embedder)?;
175
176        Ok(Self { store })
177    }
178
179    /// Get the current tokio runtime handle.
180    /// This must be called from within a tokio runtime context.
181    pub fn handle() -> Handle {
182        Handle::current()
183    }
184}