1use std::fs;
2use std::path::{Path, PathBuf};
3use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
4
5use serde::{Deserialize, Serialize};
6
7use engram_hnsw::HnswParams;
8use engram_llm_client::{
9 EmbeddingProvider, LocalTextGenerator, OpenAITextGenerator, TextGenerator,
10 VoyageEmbeddingProvider,
11};
12
13use crate::config_loader::{
14 deep_merge, load_global_config_tree, load_project_config_tree, restore_secret, secret_at,
15};
16use crate::error::CoreError;
17
18pub(crate) const ENGRAM_CONFIG_SUBPATH: &str = ".engram/engram.toml";
22const PROJECT_DIR_MARKER: &str = ".engram";
23
24const EMBEDDING_SECTION: &str = "embedding";
25const LLM_SECTION: &str = "llm";
26
27const DEFAULT_EMBEDDING_PROVIDER: &str = "voyage";
28pub const DEFAULT_EMBEDDING_MODEL: &str = "voyage-4";
29const DEFAULT_EMBEDDING_DIMENSION: usize = 1024;
30const DEFAULT_HYDE_THRESHOLD: usize = 0;
31const DEFAULT_LLM_PROVIDER: &str = "openai";
32const DEFAULT_LLM_MODEL: &str = "gpt-4o-mini";
33const DEFAULT_REINDEX_INTERVAL_SECS: u64 = 3600;
34const DEFAULT_HNSW_MAX_CONNECTIONS: usize = 16;
35const DEFAULT_HNSW_EF_CONSTRUCTION: usize = 200;
36const DEFAULT_HNSW_EF_SEARCH: usize = 40;
37const DEFAULT_CONSOLIDATION_STALE_DAYS: u32 = 90;
38const DEFAULT_CONSOLIDATION_MIN_SCORE: f64 = 0.3;
39const DEFAULT_DEDUP_THRESHOLD: f32 = 0.95;
40const DEFAULT_TRAINER_BINARY: &str = "engram-trainer";
41const DEFAULT_TRAINER_TIMEOUT_SECS: u64 = 300;
42const DEFAULT_MODELS_PATH: &str = "~/.engram/models";
43const DEFAULT_RRF_K: usize = 60;
44const DEFAULT_VECTOR_WEIGHT: f64 = 0.7;
45const DEFAULT_SPARSE_WEIGHT: f64 = 0.3;
46
47#[derive(Serialize, Deserialize, Clone)]
48pub struct Config {
49 #[serde(default)]
50 pub database: DatabaseConfig,
51 pub embedding: EmbeddingConfig,
52 pub llm: LlmConfig,
53 pub server: ServerConfig,
54 pub hnsw: HnswConfig,
55 #[serde(default)]
56 pub consolidation: ConsolidationConfig,
57 #[serde(default)]
58 pub deduplication: DeduplicationConfig,
59 #[serde(default)]
60 pub trainer: TrainerConfig,
61 #[serde(default)]
62 pub search: SearchConfig,
63}
64
65#[derive(Serialize, Deserialize, Clone, Default)]
66pub struct DatabaseConfig {
67 #[serde(default)]
72 pub path: Option<String>,
73}
74
75#[derive(Serialize, Deserialize, Clone)]
76pub struct EmbeddingConfig {
77 pub provider: String,
78 pub api_key: Option<String>,
79 pub model: Option<String>,
80 pub dimension: Option<usize>,
81 #[serde(default)]
86 pub output_dimension: Option<usize>,
87 #[serde(default)]
88 pub hyde_threshold: usize,
89}
90
91#[derive(Serialize, Deserialize, Clone)]
92pub struct LlmConfig {
93 pub provider: String,
94 pub api_key: Option<String>,
95 pub model: Option<String>,
96}
97
98#[derive(Serialize, Deserialize, Clone)]
99pub struct ServerConfig {
100 #[serde(default)]
104 pub socket_path: Option<String>,
105 pub reindex_interval_secs: u64,
106}
107
108#[derive(Serialize, Deserialize, Clone)]
109pub struct HnswConfig {
110 pub max_connections: usize,
111 pub ef_construction: usize,
112 pub ef_search: usize,
113 pub dimension: usize,
114}
115
116#[derive(Serialize, Deserialize, Clone)]
117pub struct ConsolidationConfig {
118 pub stale_days: u32,
119 pub min_score: f64,
120}
121
122#[derive(Serialize, Deserialize, Clone)]
123pub struct DeduplicationConfig {
124 pub threshold: f32,
125}
126
127#[derive(Serialize, Deserialize, Clone)]
132pub struct SearchConfig {
133 pub rrf_k: usize,
134 pub vector_weight: f64,
135 pub sparse_weight: f64,
136}
137
138#[derive(Serialize, Deserialize, Clone)]
139pub struct TrainerConfig {
140 pub trainer_binary: String,
141 pub trainer_timeout_secs: u64,
142 pub models_path: String,
143}
144
145impl Default for TrainerConfig {
146 fn default() -> Self {
147 Self {
148 trainer_binary: DEFAULT_TRAINER_BINARY.into(),
149 trainer_timeout_secs: DEFAULT_TRAINER_TIMEOUT_SECS,
150 models_path: DEFAULT_MODELS_PATH.into(),
151 }
152 }
153}
154
155impl Default for ConsolidationConfig {
156 fn default() -> Self {
157 Self {
158 stale_days: DEFAULT_CONSOLIDATION_STALE_DAYS,
159 min_score: DEFAULT_CONSOLIDATION_MIN_SCORE,
160 }
161 }
162}
163
164impl Default for DeduplicationConfig {
165 fn default() -> Self {
166 Self {
167 threshold: DEFAULT_DEDUP_THRESHOLD,
168 }
169 }
170}
171
172impl Default for SearchConfig {
173 fn default() -> Self {
174 Self {
175 rrf_k: DEFAULT_RRF_K,
176 vector_weight: DEFAULT_VECTOR_WEIGHT,
177 sparse_weight: DEFAULT_SPARSE_WEIGHT,
178 }
179 }
180}
181
182pub fn validate_dedup_threshold(threshold: f32) -> Result<(), CoreError> {
187 if threshold > 0.0 && threshold <= 1.0 {
188 return Ok(());
189 }
190 Err(CoreError::ConfigValidation(format!(
191 "deduplication.threshold must be in (0.0, 1.0], got {threshold}"
192 )))
193}
194
195pub fn validate_search_config(config: &SearchConfig) -> Result<(), CoreError> {
202 if config.rrf_k == 0 {
203 return Err(CoreError::ConfigValidation(
204 "search.rrf_k must be greater than 0".to_string(),
205 ));
206 }
207 if !config.vector_weight.is_finite() || !config.sparse_weight.is_finite() {
208 return Err(CoreError::ConfigValidation(
209 "search weights must be finite".to_string(),
210 ));
211 }
212 if config.vector_weight < 0.0 || config.sparse_weight < 0.0 {
213 return Err(CoreError::ConfigValidation(format!(
214 "search weights must be non-negative, got vector_weight={}, sparse_weight={}",
215 config.vector_weight, config.sparse_weight
216 )));
217 }
218 if config.vector_weight == 0.0 && config.sparse_weight == 0.0 {
219 return Err(CoreError::ConfigValidation(
220 "search.vector_weight and search.sparse_weight cannot both be 0.0".to_string(),
221 ));
222 }
223 Ok(())
224}
225
226impl Config {
227 pub fn load() -> Result<Self, CoreError> {
235 let global_tree = load_global_config_tree()?;
236 let project_tree = load_project_config_tree()?;
237
238 if global_tree.is_none() && project_tree.is_none() {
239 let mut config = Self::default();
240 config.apply_env_overrides();
241 return Ok(config);
242 }
243
244 let mut merged_tree = match global_tree {
245 Some(tree) => tree,
246 None => toml::Value::try_from(Self::default())
247 .map_err(|error| CoreError::ConfigParseError(error.to_string()))?,
248 };
249
250 let global_embedding_secret = secret_at(&merged_tree, EMBEDDING_SECTION);
251 let global_llm_secret = secret_at(&merged_tree, LLM_SECTION);
252
253 if let Some(project_tree) = project_tree {
254 deep_merge(&mut merged_tree, project_tree);
255 }
256
257 restore_secret(&mut merged_tree, EMBEDDING_SECTION, global_embedding_secret);
258 restore_secret(&mut merged_tree, LLM_SECTION, global_llm_secret);
259
260 let mut config: Config = merged_tree
261 .try_into()
262 .map_err(|error: toml::de::Error| CoreError::ConfigParseError(error.to_string()))?;
263 config.apply_env_overrides();
264 Ok(config)
265 }
266
267 pub fn load_from_path(path: &str) -> Result<Self, CoreError> {
268 let content = fs::read_to_string(path).map_err(|_| CoreError::ConfigNotFound)?;
269 let mut config: Config = toml::from_str(&content)
270 .map_err(|error| CoreError::ConfigParseError(error.to_string()))?;
271 config.apply_env_overrides();
272 Ok(config)
273 }
274
275 pub fn build_embedding_provider(
276 &self,
277 ) -> Result<Box<dyn EmbeddingProvider + Send + Sync>, CoreError> {
278 let dimension = self
279 .embedding
280 .dimension
281 .unwrap_or(DEFAULT_EMBEDDING_DIMENSION);
282 match self.embedding.provider.as_str() {
283 "voyage" => {
284 let api_key = self
285 .embedding
286 .api_key
287 .as_deref()
288 .filter(|key| !key.is_empty())
289 .ok_or_else(|| CoreError::InvalidProvider("voyage requires api_key".into()))?;
290 let model = self
291 .embedding
292 .model
293 .clone()
294 .unwrap_or_else(|| DEFAULT_EMBEDDING_MODEL.into());
295 let provider = VoyageEmbeddingProvider::with_config(
296 api_key.to_owned(),
297 model,
298 dimension,
299 self.embedding.output_dimension,
300 engram_llm_client::RetryConfig::default(),
301 "https://api.voyageai.com".into(),
302 )?;
303 Ok(Box::new(provider))
304 }
305 "deterministic" => Ok(Box::new(DeterministicEmbeddingProvider { dimension })),
306 other => Err(CoreError::InvalidProvider(format!(
307 "{other} embedding not supported"
308 ))),
309 }
310 }
311
312 pub fn build_text_generator(&self) -> Result<Box<dyn TextGenerator + Send + Sync>, CoreError> {
313 match self.llm.provider.as_str() {
314 "openai" => {
315 let api_key = self
316 .llm
317 .api_key
318 .as_deref()
319 .filter(|key| !key.is_empty())
320 .ok_or_else(|| CoreError::InvalidProvider("openai requires api_key".into()))?;
321 let generator = OpenAITextGenerator::new(api_key.to_owned())?;
322 Ok(Box::new(generator))
323 }
324 "local" => {
325 let models_path = expand_tilde(&self.trainer.models_path);
326 let model_path = format!("{models_path}/text_generator.onnx");
327 let tokenizer_path = format!("{models_path}/tokenizer.json");
328 let generator = LocalTextGenerator::new(&model_path, &tokenizer_path)
329 .map_err(|error| CoreError::InvalidProvider(error.to_string()))?;
330 Ok(Box::new(generator))
331 }
332 other => Err(CoreError::InvalidProvider(format!(
333 "{other} llm not supported"
334 ))),
335 }
336 }
337
338 pub fn build_hnsw_params(&self) -> Result<HnswParams, CoreError> {
339 HnswParams::new(self.hnsw.dimension)?
340 .with_max_connections(self.hnsw.max_connections)?
341 .with_ef_construction(self.hnsw.ef_construction)?
342 .with_ef_search(self.hnsw.ef_search)
343 .map_err(CoreError::Hnsw)
344 }
345
346 fn apply_env_overrides(&mut self) {
347 if let Ok(value) = std::env::var("ENGRAM_EMBEDDING_MODEL") {
352 self.embedding.model = Some(value);
353 }
354 if let Ok(value) = std::env::var("ENGRAM_LLM_MODEL") {
355 self.llm.model = Some(value);
356 }
357 if let Ok(value) = std::env::var("ENGRAM_TRAINER_BINARY") {
358 self.trainer.trainer_binary = value;
359 }
360 if let Ok(value) = std::env::var("ENGRAM_TRAINER_TIMEOUT")
361 && let Ok(secs) = value.parse::<u64>()
362 {
363 self.trainer.trainer_timeout_secs = secs;
364 }
365 if let Ok(value) = std::env::var("ENGRAM_MODELS_PATH") {
366 self.trainer.models_path = value;
367 }
368 self.apply_provider_api_key_overrides();
369 }
370
371 fn apply_provider_api_key_overrides(&mut self) {
372 if self.embedding.provider == "voyage"
373 && let Ok(value) = std::env::var("ENGRAM_VOYAGE_API_KEY")
374 {
375 self.embedding.api_key = Some(value);
376 }
377 if self.llm.provider == "openai"
378 && let Ok(value) = std::env::var("ENGRAM_OPENAI_API_KEY")
379 {
380 self.llm.api_key = Some(value);
381 }
382 }
383}
384
385impl Default for Config {
386 fn default() -> Self {
387 Self {
388 database: DatabaseConfig { path: None },
389 embedding: EmbeddingConfig {
390 provider: DEFAULT_EMBEDDING_PROVIDER.into(),
391 api_key: None,
392 model: Some(DEFAULT_EMBEDDING_MODEL.into()),
393 dimension: Some(DEFAULT_EMBEDDING_DIMENSION),
394 output_dimension: None,
395 hyde_threshold: DEFAULT_HYDE_THRESHOLD,
396 },
397 llm: LlmConfig {
398 provider: DEFAULT_LLM_PROVIDER.into(),
399 api_key: None,
400 model: Some(DEFAULT_LLM_MODEL.into()),
401 },
402 server: ServerConfig {
403 socket_path: None,
404 reindex_interval_secs: DEFAULT_REINDEX_INTERVAL_SECS,
405 },
406 hnsw: HnswConfig {
407 max_connections: DEFAULT_HNSW_MAX_CONNECTIONS,
408 ef_construction: DEFAULT_HNSW_EF_CONSTRUCTION,
409 ef_search: DEFAULT_HNSW_EF_SEARCH,
410 dimension: DEFAULT_EMBEDDING_DIMENSION,
411 },
412 consolidation: ConsolidationConfig::default(),
413 deduplication: DeduplicationConfig::default(),
414 trainer: TrainerConfig::default(),
415 search: SearchConfig::default(),
416 }
417 }
418}
419
420pub fn resolve_project_dir(
421 start: &Path,
422 explicit_override: Option<&Path>,
423) -> Result<PathBuf, CoreError> {
424 if let Some(path) = explicit_override {
425 return Ok(path.to_path_buf());
426 }
427 if let Ok(env_path) = std::env::var("ENGRAM_PROJECT_DIR") {
428 let candidate = PathBuf::from(env_path);
429 if candidate.is_absolute() {
430 return Ok(candidate);
431 }
432 }
433 let mut current = start.to_path_buf();
434 loop {
435 if current.join(PROJECT_DIR_MARKER).is_dir() {
436 return Ok(current);
437 }
438 if !current.pop() {
440 break;
441 }
442 }
443 Err(CoreError::ProjectDirNotFound)
444}
445
446pub fn expand_tilde(path: &str) -> String {
447 if let Some(rest) = path.strip_prefix("~/")
448 && let Some(home) = home_directory()
449 {
450 return format!("{home}/{rest}");
451 }
452 path.to_string()
453}
454
455pub fn home_directory() -> Option<String> {
456 std::env::var("HOME").ok()
457}
458
459struct DeterministicEmbeddingProvider {
460 dimension: usize,
461}
462
463static DETERMINISTIC_PROVIDER_INSTRUMENTATION: AtomicBool = AtomicBool::new(false);
470static DETERMINISTIC_PROVIDER_ENTRIES: AtomicUsize = AtomicUsize::new(0);
471static DETERMINISTIC_PROVIDER_MAX_CONCURRENT: AtomicUsize = AtomicUsize::new(0);
472
473pub fn enable_deterministic_provider_instrumentation() {
477 DETERMINISTIC_PROVIDER_INSTRUMENTATION.store(true, Ordering::Relaxed);
478}
479
480pub fn disable_deterministic_provider_instrumentation() {
483 DETERMINISTIC_PROVIDER_INSTRUMENTATION.store(false, Ordering::Relaxed);
484}
485
486pub fn deterministic_provider_max_concurrent() -> usize {
487 DETERMINISTIC_PROVIDER_MAX_CONCURRENT.load(Ordering::Relaxed)
488}
489
490pub fn reset_deterministic_provider_counters() {
491 DETERMINISTIC_PROVIDER_ENTRIES.store(0, Ordering::Relaxed);
492 DETERMINISTIC_PROVIDER_MAX_CONCURRENT.store(0, Ordering::Relaxed);
493}
494
495impl EmbeddingProvider for DeterministicEmbeddingProvider {
496 fn embed(
497 &self,
498 text: &str,
499 _input_type: Option<&str>,
500 ) -> Result<Vec<f32>, engram_llm_client::ApiError> {
501 let instrumented = DETERMINISTIC_PROVIDER_INSTRUMENTATION.load(Ordering::Relaxed);
508 if instrumented {
509 let entries = DETERMINISTIC_PROVIDER_ENTRIES.fetch_add(1, Ordering::SeqCst) + 1;
510 DETERMINISTIC_PROVIDER_MAX_CONCURRENT.fetch_max(entries, Ordering::SeqCst);
511 std::thread::sleep(std::time::Duration::from_millis(20));
512 }
513
514 let mut embedding = vec![0.0_f32; self.dimension];
515 for (index, byte) in text.bytes().enumerate() {
516 embedding[index % self.dimension] += byte as f32 / 255.0;
517 }
518 let norm: f32 = embedding
519 .iter()
520 .map(|v| v * v)
521 .sum::<f32>()
522 .sqrt()
523 .max(1e-10);
524 for value in &mut embedding {
525 *value /= norm;
526 }
527
528 if instrumented {
529 DETERMINISTIC_PROVIDER_ENTRIES.fetch_sub(1, Ordering::SeqCst);
530 }
531 Ok(embedding)
532 }
533
534 fn dimension(&self) -> usize {
535 self.dimension
536 }
537
538 fn model_name(&self) -> &str {
539 "deterministic"
540 }
541}
542
543#[cfg(test)]
544mod tests {
545 use super::*;
546
547 #[test]
548 fn deterministic_provider_ignores_input_type() {
549 let provider = DeterministicEmbeddingProvider { dimension: 8 };
550 let text = "rust ownership rules govern compile-time memory access";
551
552 let document = provider.embed(text, Some("document")).unwrap();
553 let query = provider.embed(text, Some("query")).unwrap();
554 let none = provider.embed(text, None).unwrap();
555
556 assert_eq!(document, query);
559 assert_eq!(document, none);
560 }
561
562 #[test]
563 fn embedding_config_deserializes_without_output_dimension() {
564 let toml_input = r#"
565 provider = "voyage"
566 model = "voyage-code-3"
567 dimension = 1024
568 "#;
569 let config: EmbeddingConfig = toml::from_str(toml_input).unwrap();
570 assert_eq!(config.provider, "voyage");
571 assert_eq!(config.output_dimension, None);
572 }
573
574 #[test]
575 fn embedding_config_deserializes_with_output_dimension() {
576 let toml_input = r#"
577 provider = "voyage"
578 model = "voyage-4"
579 dimension = 1024
580 output_dimension = 1024
581 "#;
582 let config: EmbeddingConfig = toml::from_str(toml_input).unwrap();
583 assert_eq!(config.output_dimension, Some(1024));
584 }
585
586 #[test]
587 fn default_dedup_threshold_is_documented_default() {
588 assert_eq!(Config::default().deduplication.threshold, 0.95);
589 }
590
591 #[test]
592 fn validate_dedup_threshold_rejects_zero() {
593 let error = validate_dedup_threshold(0.0).expect_err("zero is out of range");
594 assert!(matches!(error, CoreError::ConfigValidation(_)));
595 assert!(error.to_string().contains("[6022]"));
596 assert!(error.to_string().contains("0"));
597 }
598
599 #[test]
600 fn validate_dedup_threshold_rejects_negative() {
601 let error = validate_dedup_threshold(-0.5).expect_err("negative is out of range");
602 assert!(matches!(error, CoreError::ConfigValidation(_)));
603 assert!(error.to_string().contains("-0.5"));
604 }
605
606 #[test]
607 fn validate_dedup_threshold_rejects_above_one() {
608 let error = validate_dedup_threshold(1.0001).expect_err("above 1.0 is out of range");
609 assert!(matches!(error, CoreError::ConfigValidation(_)));
610 }
611
612 #[test]
613 fn validate_dedup_threshold_accepts_one() {
614 validate_dedup_threshold(1.0).expect("1.0 is the inclusive upper bound");
615 }
616
617 #[test]
618 fn validate_dedup_threshold_accepts_tiny_positive() {
619 validate_dedup_threshold(0.0001).expect("any positive within range is accepted");
620 }
621
622 #[test]
623 fn default_search_config_matches_documented_values() {
624 let search = SearchConfig::default();
625 assert_eq!(search.rrf_k, 60);
626 assert_eq!(search.vector_weight, 0.7);
627 assert_eq!(search.sparse_weight, 0.3);
628 }
629
630 #[test]
631 fn config_loads_without_search_section() {
632 let toml_input = r#"
635 [embedding]
636 provider = "deterministic"
637
638 [llm]
639 provider = "openai"
640
641 [server]
642 reindex_interval_secs = 3600
643
644 [hnsw]
645 max_connections = 16
646 ef_construction = 200
647 ef_search = 40
648 dimension = 1024
649 "#;
650 let config: Config = toml::from_str(toml_input).unwrap();
651 assert_eq!(config.search.rrf_k, 60);
652 assert_eq!(config.search.vector_weight, 0.7);
653 assert_eq!(config.search.sparse_weight, 0.3);
654 }
655
656 #[test]
657 fn search_config_round_trips_through_toml() {
658 let toml_input = r#"
659 rrf_k = 30
660 vector_weight = 0.6
661 sparse_weight = 0.4
662 "#;
663 let search: SearchConfig = toml::from_str(toml_input).unwrap();
664 assert_eq!(search.rrf_k, 30);
665 assert_eq!(search.vector_weight, 0.6);
666 assert_eq!(search.sparse_weight, 0.4);
667 }
668
669 #[test]
670 fn validate_search_config_accepts_defaults() {
671 validate_search_config(&SearchConfig::default()).expect("documented defaults are valid");
672 }
673
674 #[test]
675 fn validate_search_config_accepts_one_zero_weight() {
676 let search = SearchConfig {
677 rrf_k: 60,
678 vector_weight: 0.0,
679 sparse_weight: 0.3,
680 };
681 validate_search_config(&search).expect("a single zero weight still leaves one live term");
682 }
683
684 #[test]
685 fn validate_search_config_rejects_zero_rrf_k() {
686 let search = SearchConfig {
687 rrf_k: 0,
688 ..SearchConfig::default()
689 };
690 let error = validate_search_config(&search).expect_err("rrf_k of 0 divides by k+rank base");
691 assert!(matches!(error, CoreError::ConfigValidation(_)));
692 assert!(error.to_string().contains("rrf_k"));
693 }
694
695 #[test]
696 fn validate_search_config_rejects_negative_weight() {
697 let search = SearchConfig {
698 rrf_k: 60,
699 vector_weight: -0.1,
700 sparse_weight: 0.3,
701 };
702 let error = validate_search_config(&search).expect_err("negative weight is invalid");
703 assert!(matches!(error, CoreError::ConfigValidation(_)));
704 assert!(error.to_string().contains("non-negative"));
705 }
706
707 #[test]
708 fn validate_search_config_rejects_non_finite_weight() {
709 let search = SearchConfig {
710 rrf_k: 60,
711 vector_weight: f64::NAN,
712 sparse_weight: 0.3,
713 };
714 let error = validate_search_config(&search).expect_err("NaN weight is invalid");
715 assert!(matches!(error, CoreError::ConfigValidation(_)));
716 assert!(error.to_string().contains("finite"));
717 }
718
719 #[test]
720 fn validate_search_config_rejects_all_zero_weights() {
721 let search = SearchConfig {
722 rrf_k: 60,
723 vector_weight: 0.0,
724 sparse_weight: 0.0,
725 };
726 let error = validate_search_config(&search).expect_err("both weights zero cancels output");
727 assert!(matches!(error, CoreError::ConfigValidation(_)));
728 }
729}