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

engram_core/
config.rs
