semantic_memory_mcp/
bridge.rs1#[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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub enum EmbedderBackend {
17 Candle,
20 Ollama,
22 Mock,
24}
25
26impl Default for EmbedderBackend {
27 fn default() -> Self {
28 #[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 pub memory_dir: PathBuf,
66 pub embedder_backend: EmbedderBackend,
67 pub embedding_url: String,
69 pub embedding_model: String,
71 pub embedding_dims: usize,
72 pub turbo_quant_enabled: bool,
74 pub turbo_quant_bits: Option<u8>,
76 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 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 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 #[cfg(feature = "full")]
153 {
154 if config.turbo_quant_enabled {
155 use semantic_memory::DerivedVectorBackendPolicy;
156 search_config.derived_vector_backend = DerivedVectorBackendPolicy::TurboQuantCandidateOnly;
157 if let Some(bits) = config.turbo_quant_bits {
158 search_config.turbo_quant_bits = bits;
159 }
160 if let Some(projs) = config.turbo_quant_projections {
161 search_config.turbo_quant_projections = projs;
162 }
163 }
164 }
165
166 let mem_config = MemoryConfig {
167 base_dir: config.memory_dir,
168 embedding: embedding_config,
169 search: search_config,
170 ..Default::default()
171 };
172
173 let store = MemoryStore::open_with_embedder(mem_config, embedder)?;
174
175 Ok(Self { store })
176 }
177
178 pub fn handle() -> Handle {
181 Handle::current()
182 }
183}