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, Default)]
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    #[cfg_attr(feature = "candle-embedder", default)]
20    Candle,
21    /// External Ollama server (requires `ollama serve` running).
22    #[cfg_attr(not(feature = "candle-embedder"), default)]
23    Ollama,
24    /// Mock embedder for testing (deterministic hash-based, no real embeddings).
25    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    /// MCP-005: renamed from db_path to memory_dir — this is a directory,
52    /// not a SQLite file path. semantic-memory creates base_dir/memory.db
53    /// inside it.
54    pub memory_dir: PathBuf,
55    pub embedder_backend: EmbedderBackend,
56    /// Ollama URL — only used when backend is Ollama.
57    pub embedding_url: String,
58    /// Embedding model name — used by both Candle (as HF model ID) and Ollama.
59    pub embedding_model: String,
60    pub embedding_dims: usize,
61    /// Enable TurboQuant compressed vector candidate backend.
62    pub turbo_quant_enabled: bool,
63    /// TurboQuant polar angle bits (default: 8).
64    pub turbo_quant_bits: Option<u8>,
65    /// TurboQuant QJL projection count (default: 16).
66    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    /// Open the memory store with the given config.
98    /// Store opening is synchronous — no runtime needed here.
99    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                    // For Candle, the model field is the HuggingFace model ID.
115                    // Map common Ollama model names to HF model IDs.
116                    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        // TurboQuant compressed vector candidate backend
143        #[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    /// Get the current tokio runtime handle.
172    /// This must be called from within a tokio runtime context.
173    pub fn handle() -> Handle {
174        Handle::current()
175    }
176}