engram_core/
config.rs

1use std::fs;
2use std::path::{Path, PathBuf};
3use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
4
5use serde::Deserialize;
6
7use engram_hnsw::HnswParams;
8use engram_llm_client::{
9    EmbeddingProvider, LocalTextGenerator, OpenAITextGenerator, TextGenerator,
10    VoyageEmbeddingProvider,
11};
12
13use crate::error::CoreError;
14
15const CONFIG_HOME_SUBDIR: &str = ".engram/engram.toml";
16const PROJECT_DIR_MARKER: &str = ".engram";
17
18const DEFAULT_DB_PATH: &str = "~/.engram/memories.db";
19const DEFAULT_SOCKET_PATH: &str = "~/.engram/engram.sock";
20const DEFAULT_EMBEDDING_PROVIDER: &str = "voyage";
21const DEFAULT_EMBEDDING_MODEL: &str = "voyage-code-3";
22const DEFAULT_EMBEDDING_DIMENSION: usize = 1024;
23const DEFAULT_HYDE_THRESHOLD: usize = 0;
24const DEFAULT_LLM_PROVIDER: &str = "openai";
25const DEFAULT_LLM_MODEL: &str = "gpt-4o-mini";
26const DEFAULT_REINDEX_INTERVAL_SECS: u64 = 3600;
27const DEFAULT_HNSW_MAX_CONNECTIONS: usize = 16;
28const DEFAULT_HNSW_EF_CONSTRUCTION: usize = 200;
29const DEFAULT_HNSW_EF_SEARCH: usize = 40;
30const DEFAULT_CONSOLIDATION_STALE_DAYS: u32 = 90;
31const DEFAULT_CONSOLIDATION_MIN_SCORE: f64 = 0.3;
32const DEFAULT_TRAINER_BINARY: &str = "engram-trainer";
33const DEFAULT_TRAINER_TIMEOUT_SECS: u64 = 300;
34const DEFAULT_MODELS_PATH: &str = "~/.engram/models";
35
36#[derive(Deserialize, Clone)]
37pub struct Config {
38    pub database: DatabaseConfig,
39    pub embedding: EmbeddingConfig,
40    pub llm: LlmConfig,
41    pub server: ServerConfig,
42    pub hnsw: HnswConfig,
43    #[serde(default)]
44    pub consolidation: ConsolidationConfig,
45    #[serde(default)]
46    pub trainer: TrainerConfig,
47}
48
49#[derive(Deserialize, Clone)]
50pub struct DatabaseConfig {
51    pub path: String,
52}
53
54#[derive(Deserialize, Clone)]
55pub struct EmbeddingConfig {
56    pub provider: String,
57    pub api_key: Option<String>,
58    pub model: Option<String>,
59    pub dimension: Option<usize>,
60    #[serde(default)]
61    pub hyde_threshold: usize,
62}
63
64#[derive(Deserialize, Clone)]
65pub struct LlmConfig {
66    pub provider: String,
67    pub api_key: Option<String>,
68    pub model: Option<String>,
69}
70
71#[derive(Deserialize, Clone)]
72pub struct ServerConfig {
73    pub socket_path: String,
74    pub reindex_interval_secs: u64,
75}
76
77#[derive(Deserialize, Clone)]
78pub struct HnswConfig {
79    pub max_connections: usize,
80    pub ef_construction: usize,
81    pub ef_search: usize,
82    pub dimension: usize,
83}
84
85#[derive(Deserialize, Clone)]
86pub struct ConsolidationConfig {
87    pub stale_days: u32,
88    pub min_score: f64,
89}
90
91#[derive(Deserialize, Clone)]
92pub struct TrainerConfig {
93    pub trainer_binary: String,
94    pub trainer_timeout_secs: u64,
95    pub models_path: String,
96}
97
98impl Default for TrainerConfig {
99    fn default() -> Self {
100        Self {
101            trainer_binary: DEFAULT_TRAINER_BINARY.into(),
102            trainer_timeout_secs: DEFAULT_TRAINER_TIMEOUT_SECS,
103            models_path: DEFAULT_MODELS_PATH.into(),
104        }
105    }
106}
107
108impl Default for ConsolidationConfig {
109    fn default() -> Self {
110        Self {
111            stale_days: DEFAULT_CONSOLIDATION_STALE_DAYS,
112            min_score: DEFAULT_CONSOLIDATION_MIN_SCORE,
113        }
114    }
115}
116
117impl Config {
118    pub fn load() -> Result<Self, CoreError> {
119        if let Some(home) = home_directory() {
120            let home_config = Path::new(&home).join(CONFIG_HOME_SUBDIR);
121            if home_config.exists() {
122                return Self::load_from_path(home_config.to_str().unwrap_or_default());
123            }
124        }
125
126        let mut config = Self::default();
127        config.apply_env_overrides();
128        Ok(config)
129    }
130
131    pub fn load_from_path(path: &str) -> Result<Self, CoreError> {
132        let content = fs::read_to_string(path).map_err(|_| CoreError::ConfigNotFound)?;
133        let mut config: Config = toml::from_str(&content)
134            .map_err(|error| CoreError::ConfigParseError(error.to_string()))?;
135        config.apply_env_overrides();
136        Ok(config)
137    }
138
139    pub fn build_embedding_provider(
140        &self,
141    ) -> Result<Box<dyn EmbeddingProvider + Send + Sync>, CoreError> {
142        let dimension = self
143            .embedding
144            .dimension
145            .unwrap_or(DEFAULT_EMBEDDING_DIMENSION);
146        match self.embedding.provider.as_str() {
147            "voyage" => {
148                let api_key = self
149                    .embedding
150                    .api_key
151                    .as_deref()
152                    .filter(|key| !key.is_empty())
153                    .ok_or_else(|| CoreError::InvalidProvider("voyage requires api_key".into()))?;
154                let provider = VoyageEmbeddingProvider::new(api_key.to_owned())?;
155                Ok(Box::new(provider))
156            }
157            "deterministic" => Ok(Box::new(DeterministicEmbeddingProvider { dimension })),
158            other => Err(CoreError::InvalidProvider(format!(
159                "{other} embedding not supported"
160            ))),
161        }
162    }
163
164    pub fn build_text_generator(&self) -> Result<Box<dyn TextGenerator + Send + Sync>, CoreError> {
165        match self.llm.provider.as_str() {
166            "openai" => {
167                let api_key = self
168                    .llm
169                    .api_key
170                    .as_deref()
171                    .filter(|key| !key.is_empty())
172                    .ok_or_else(|| CoreError::InvalidProvider("openai requires api_key".into()))?;
173                let generator = OpenAITextGenerator::new(api_key.to_owned())?;
174                Ok(Box::new(generator))
175            }
176            "local" => {
177                let models_path = expand_tilde(&self.trainer.models_path);
178                let model_path = format!("{models_path}/text_generator.onnx");
179                let tokenizer_path = format!("{models_path}/tokenizer.json");
180                let generator = LocalTextGenerator::new(&model_path, &tokenizer_path)
181                    .map_err(|error| CoreError::InvalidProvider(error.to_string()))?;
182                Ok(Box::new(generator))
183            }
184            other => Err(CoreError::InvalidProvider(format!(
185                "{other} llm not supported"
186            ))),
187        }
188    }
189
190    pub fn build_hnsw_params(&self) -> Result<HnswParams, CoreError> {
191        HnswParams::new(self.hnsw.dimension)?
192            .with_max_connections(self.hnsw.max_connections)?
193            .with_ef_construction(self.hnsw.ef_construction)?
194            .with_ef_search(self.hnsw.ef_search)
195            .map_err(CoreError::Hnsw)
196    }
197
198    fn apply_env_overrides(&mut self) {
199        // ENGRAM_DB_PATH / ENGRAM_SOCKET_PATH are honored directly by server.rs
200        // (resolve_database_path / resolve_socket_path). Copying them onto
201        // `self.database.path` / `self.server.socket_path` here would be redundant
202        // and could mask the per-project layout derived from `project_dir`.
203        if let Ok(value) = std::env::var("ENGRAM_EMBEDDING_MODEL") {
204            self.embedding.model = Some(value);
205        }
206        if let Ok(value) = std::env::var("ENGRAM_LLM_MODEL") {
207            self.llm.model = Some(value);
208        }
209        if let Ok(value) = std::env::var("ENGRAM_TRAINER_BINARY") {
210            self.trainer.trainer_binary = value;
211        }
212        if let Ok(value) = std::env::var("ENGRAM_TRAINER_TIMEOUT")
213            && let Ok(secs) = value.parse::<u64>()
214        {
215            self.trainer.trainer_timeout_secs = secs;
216        }
217        if let Ok(value) = std::env::var("ENGRAM_MODELS_PATH") {
218            self.trainer.models_path = value;
219        }
220        self.apply_provider_api_key_overrides();
221    }
222
223    fn apply_provider_api_key_overrides(&mut self) {
224        if self.embedding.provider == "voyage"
225            && let Ok(value) = std::env::var("ENGRAM_VOYAGE_API_KEY")
226        {
227            self.embedding.api_key = Some(value);
228        }
229        if self.llm.provider == "openai"
230            && let Ok(value) = std::env::var("ENGRAM_OPENAI_API_KEY")
231        {
232            self.llm.api_key = Some(value);
233        }
234    }
235}
236
237impl Default for Config {
238    fn default() -> Self {
239        Self {
240            database: DatabaseConfig {
241                path: DEFAULT_DB_PATH.into(),
242            },
243            embedding: EmbeddingConfig {
244                provider: DEFAULT_EMBEDDING_PROVIDER.into(),
245                api_key: None,
246                model: Some(DEFAULT_EMBEDDING_MODEL.into()),
247                dimension: Some(DEFAULT_EMBEDDING_DIMENSION),
248                hyde_threshold: DEFAULT_HYDE_THRESHOLD,
249            },
250            llm: LlmConfig {
251                provider: DEFAULT_LLM_PROVIDER.into(),
252                api_key: None,
253                model: Some(DEFAULT_LLM_MODEL.into()),
254            },
255            server: ServerConfig {
256                socket_path: DEFAULT_SOCKET_PATH.into(),
257                reindex_interval_secs: DEFAULT_REINDEX_INTERVAL_SECS,
258            },
259            hnsw: HnswConfig {
260                max_connections: DEFAULT_HNSW_MAX_CONNECTIONS,
261                ef_construction: DEFAULT_HNSW_EF_CONSTRUCTION,
262                ef_search: DEFAULT_HNSW_EF_SEARCH,
263                dimension: DEFAULT_EMBEDDING_DIMENSION,
264            },
265            consolidation: ConsolidationConfig::default(),
266            trainer: TrainerConfig::default(),
267        }
268    }
269}
270
271pub fn resolve_project_dir(
272    start: &Path,
273    explicit_override: Option<&Path>,
274) -> Result<PathBuf, CoreError> {
275    if let Some(path) = explicit_override {
276        return Ok(path.to_path_buf());
277    }
278    if let Ok(env_path) = std::env::var("ENGRAM_PROJECT_DIR") {
279        let candidate = PathBuf::from(env_path);
280        if candidate.is_absolute() {
281            return Ok(candidate);
282        }
283    }
284    let mut current = start.to_path_buf();
285    loop {
286        if current.join(PROJECT_DIR_MARKER).is_dir() {
287            return Ok(current);
288        }
289        // `PathBuf::pop()` returns false at filesystem root, which terminates this loop safely.
290        if !current.pop() {
291            break;
292        }
293    }
294    Err(CoreError::ProjectDirNotFound)
295}
296
297pub fn expand_tilde(path: &str) -> String {
298    if let Some(rest) = path.strip_prefix("~/")
299        && let Some(home) = home_directory()
300    {
301        return format!("{home}/{rest}");
302    }
303    path.to_string()
304}
305
306pub fn home_directory() -> Option<String> {
307    std::env::var("HOME").ok()
308}
309
310struct DeterministicEmbeddingProvider {
311    dimension: usize,
312}
313
314// Test instrumentation for the deterministic provider. Gated behind
315// `DETERMINISTIC_PROVIDER_INSTRUMENTATION`, which defaults to OFF so
316// production users (TUI demo mode, local setups without API keys) pay
317// zero overhead. Tests call `enable_deterministic_provider_instrumentation`
318// before parallel assertions and `disable_deterministic_provider_instrumentation`
319// in teardown (typically via a drop-guard).
320static DETERMINISTIC_PROVIDER_INSTRUMENTATION: AtomicBool = AtomicBool::new(false);
321static DETERMINISTIC_PROVIDER_ENTRIES: AtomicUsize = AtomicUsize::new(0);
322static DETERMINISTIC_PROVIDER_MAX_CONCURRENT: AtomicUsize = AtomicUsize::new(0);
323
324/// Enable test-only instrumentation (counter + 20ms sleep) on the deterministic
325/// embedding provider. Off in production by default; tests enable it before
326/// reading concurrency counters.
327pub fn enable_deterministic_provider_instrumentation() {
328    DETERMINISTIC_PROVIDER_INSTRUMENTATION.store(true, Ordering::Relaxed);
329}
330
331/// Disable instrumentation. Tests should call this in teardown to avoid leaking
332/// state across serially-run suites.
333pub fn disable_deterministic_provider_instrumentation() {
334    DETERMINISTIC_PROVIDER_INSTRUMENTATION.store(false, Ordering::Relaxed);
335}
336
337pub fn deterministic_provider_max_concurrent() -> usize {
338    DETERMINISTIC_PROVIDER_MAX_CONCURRENT.load(Ordering::Relaxed)
339}
340
341pub fn reset_deterministic_provider_counters() {
342    DETERMINISTIC_PROVIDER_ENTRIES.store(0, Ordering::Relaxed);
343    DETERMINISTIC_PROVIDER_MAX_CONCURRENT.store(0, Ordering::Relaxed);
344}
345
346impl EmbeddingProvider for DeterministicEmbeddingProvider {
347    fn embed(&self, text: &str) -> Result<Vec<f32>, engram_llm_client::ApiError> {
348        // Test-only instrumentation: when enabled via
349        // `enable_deterministic_provider_instrumentation`, each embed call
350        // records concurrent entries and sleeps 20ms to widen the overlap
351        // window for parallelism assertions. Production users never enable
352        // this — a relaxed atomic load is ~1ns, so overhead is effectively
353        // zero otherwise.
354        let instrumented = DETERMINISTIC_PROVIDER_INSTRUMENTATION.load(Ordering::Relaxed);
355        if instrumented {
356            let entries = DETERMINISTIC_PROVIDER_ENTRIES.fetch_add(1, Ordering::SeqCst) + 1;
357            DETERMINISTIC_PROVIDER_MAX_CONCURRENT.fetch_max(entries, Ordering::SeqCst);
358            std::thread::sleep(std::time::Duration::from_millis(20));
359        }
360
361        let mut embedding = vec![0.0_f32; self.dimension];
362        for (index, byte) in text.bytes().enumerate() {
363            embedding[index % self.dimension] += byte as f32 / 255.0;
364        }
365        let norm: f32 = embedding
366            .iter()
367            .map(|v| v * v)
368            .sum::<f32>()
369            .sqrt()
370            .max(1e-10);
371        for value in &mut embedding {
372            *value /= norm;
373        }
374
375        if instrumented {
376            DETERMINISTIC_PROVIDER_ENTRIES.fetch_sub(1, Ordering::SeqCst);
377        }
378        Ok(embedding)
379    }
380
381    fn dimension(&self) -> usize {
382        self.dimension
383    }
384
385    fn model_name(&self) -> &str {
386        "deterministic"
387    }
388}
engram_core/config.rs

engram_core/
config.rs
