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, Default)]
16pub enum EmbedderBackend {
17 #[cfg_attr(feature = "candle-embedder", default)]
20 Candle,
21 #[cfg_attr(not(feature = "candle-embedder"), default)]
23 Ollama,
24 Mock,
26}
27
28impl std::str::FromStr for EmbedderBackend {
29 type Err = String;
30
31 fn from_str(s: &str) -> Result<Self, Self::Err> {
32 match s.to_lowercase().as_str() {
33 "candle" => Ok(EmbedderBackend::Candle),
34 "ollama" => Ok(EmbedderBackend::Ollama),
35 "mock" => Ok(EmbedderBackend::Mock),
36 other => Err(format!(
37 "unknown embedder '{other}', expected: candle, ollama, or mock"
38 )),
39 }
40 }
41}
42
43#[derive(Clone)]
44pub struct MemoryBridge {
45 pub store: MemoryStore,
46 pub memory_dir: PathBuf,
47}
48
49#[derive(Clone)]
50pub struct BridgeConfig {
51 pub memory_dir: PathBuf,
55 pub embedder_backend: EmbedderBackend,
56 pub embedding_url: String,
58 pub embedding_model: String,
60 pub embedding_dims: usize,
61 pub turbo_quant_enabled: bool,
63 pub turbo_quant_bits: Option<u8>,
65 pub turbo_quant_projections: Option<usize>,
67}
68
69impl BridgeConfig {
70 #[allow(clippy::too_many_arguments)]
71 pub fn from_args(
72 memory_dir: &str,
73 embedder_backend: EmbedderBackend,
74 embedding_url: Option<&str>,
75 embedding_model: Option<&str>,
76 embedding_dims: Option<usize>,
77 turbo_quant_enabled: bool,
78 turbo_quant_bits: Option<u8>,
79 turbo_quant_projections: Option<usize>,
80 ) -> Self {
81 Self {
82 memory_dir: PathBuf::from(memory_dir),
83 embedder_backend,
84 embedding_url: embedding_url
85 .unwrap_or("http://localhost:11434")
86 .to_string(),
87 embedding_model: embedding_model.unwrap_or("nomic-embed-text").to_string(),
88 embedding_dims: embedding_dims.unwrap_or(768),
89 turbo_quant_enabled,
90 turbo_quant_bits,
91 turbo_quant_projections,
92 }
93 }
94}
95
96impl MemoryBridge {
97 pub fn open(config: BridgeConfig) -> anyhow::Result<Self> {
100 let embedding_config = EmbeddingConfig {
101 ollama_url: config.embedding_url,
102 model: config.embedding_model,
103 dimensions: config.embedding_dims,
104 batch_size: 32,
105 timeout_secs: 60,
106 };
107
108 let embedder: Box<dyn Embedder> = match config.embedder_backend {
109 EmbedderBackend::Mock => Box::new(MockEmbedder::new(config.embedding_dims)),
110 EmbedderBackend::Ollama => Box::new(OllamaEmbedder::try_new(&embedding_config)?),
111 EmbedderBackend::Candle => {
112 #[cfg(feature = "candle-embedder")]
113 {
114 let hf_model_id = match embedding_config.model.as_str() {
117 "nomic-embed-text" | "nomic-embed-text-v1.5" => {
118 "nomic-ai/nomic-embed-text-v1.5"
119 }
120 other => other,
121 };
122 let candle_config = EmbeddingConfig {
123 model: hf_model_id.to_string(),
124 ..embedding_config.clone()
125 };
126 Box::new(CandleEmbedder::try_new(&candle_config)?)
127 }
128 #[cfg(not(feature = "candle-embedder"))]
129 {
130 anyhow::bail!(
131 "candle embedder requested but the 'candle-embedder' feature \
132 is not enabled. Rebuild with --features candle-embedder \
133 or use --embedder ollama"
134 )
135 }
136 }
137 };
138
139 #[allow(unused_mut)]
140 let mut search_config = SearchConfig::default();
141
142 #[cfg(feature = "full")]
144 {
145 if config.turbo_quant_enabled {
146 use semantic_memory::DerivedVectorBackendPolicy;
147 search_config.derived_vector_backend =
148 DerivedVectorBackendPolicy::TurboQuantCandidateOnly;
149 if let Some(bits) = config.turbo_quant_bits {
150 search_config.turbo_quant_bits = bits;
151 }
152 if let Some(projs) = config.turbo_quant_projections {
153 search_config.turbo_quant_projections = projs;
154 }
155 }
156 }
157
158 let memory_dir = config.memory_dir.clone();
159 let mem_config = MemoryConfig {
160 base_dir: config.memory_dir,
161 embedding: embedding_config,
162 search: search_config,
163 ..Default::default()
164 };
165
166 let store = MemoryStore::open_with_embedder(mem_config, embedder)?;
167
168 Ok(Self { store, memory_dir })
169 }
170
171 pub fn handle() -> Handle {
174 Handle::current()
175 }
176}