use std::path::PathBuf;
use std::time::Duration;
use serde::{Deserialize, Serialize};
use crate::error::ValidationError;
use crate::types::CollectiveId;
#[derive(Clone, Debug)]
pub struct Config {
pub embedding_provider: EmbeddingProvider,
pub embedding_dimension: EmbeddingDimension,
pub default_collective: Option<CollectiveId>,
pub cache_size_mb: usize,
pub sync_mode: SyncMode,
pub hnsw: HnswConfig,
pub activity: ActivityConfig,
pub watch: WatchConfig,
pub decay: DecayConfig,
pub read_only: bool,
pub migration_available_memory_bytes: Option<u64>,
}
impl Default for Config {
fn default() -> Self {
Self {
embedding_provider: EmbeddingProvider::External,
embedding_dimension: EmbeddingDimension::D384,
default_collective: None,
cache_size_mb: 64,
sync_mode: SyncMode::Normal,
hnsw: HnswConfig::default(),
activity: ActivityConfig::default(),
watch: WatchConfig::default(),
decay: DecayConfig::default(),
read_only: false,
migration_available_memory_bytes: None,
}
}
}
impl Config {
pub fn new() -> Self {
Self::default()
}
pub fn read_only() -> Self {
Self {
read_only: true,
..Default::default()
}
}
pub fn with_builtin_embeddings() -> Self {
Self {
embedding_provider: EmbeddingProvider::Builtin { model_path: None },
..Default::default()
}
}
pub fn with_external_embeddings(dimension: EmbeddingDimension) -> Self {
Self {
embedding_provider: EmbeddingProvider::External,
embedding_dimension: dimension,
..Default::default()
}
}
pub fn validate(&self) -> Result<(), ValidationError> {
if self.cache_size_mb == 0 {
return Err(ValidationError::invalid_field(
"cache_size_mb",
"must be greater than 0",
));
}
if self.hnsw.max_nb_connection == 0 {
return Err(ValidationError::invalid_field(
"hnsw.max_nb_connection",
"must be greater than 0",
));
}
if self.hnsw.ef_construction == 0 {
return Err(ValidationError::invalid_field(
"hnsw.ef_construction",
"must be greater than 0",
));
}
if self.hnsw.ef_search == 0 {
return Err(ValidationError::invalid_field(
"hnsw.ef_search",
"must be greater than 0",
));
}
if self.watch.buffer_size == 0 {
return Err(ValidationError::invalid_field(
"watch.buffer_size",
"must be greater than 0",
));
}
if self.watch.poll_interval_ms == 0 {
return Err(ValidationError::invalid_field(
"watch.poll_interval_ms",
"must be greater than 0",
));
}
if self.decay.half_life.is_zero() {
return Err(ValidationError::invalid_field(
"decay.half_life",
"must be greater than 0",
));
}
if !self.decay.freq_weight.is_finite() || self.decay.freq_weight < 0.0 {
return Err(ValidationError::invalid_field(
"decay.freq_weight",
"must be finite and non-negative",
));
}
if !self.decay.floor.is_finite() || !(0.0..=1.0).contains(&self.decay.floor) {
return Err(ValidationError::invalid_field(
"decay.floor",
"must be finite and between 0 and 1",
));
}
if let Some(weights) = self.decay.default_recall_weights {
weights.validate("decay.default_recall_weights")?;
}
if let EmbeddingDimension::Custom(dim) = self.embedding_dimension {
if dim == 0 {
return Err(ValidationError::invalid_field(
"embedding_dimension",
"custom dimension must be greater than 0",
));
}
if dim > 4096 {
return Err(ValidationError::invalid_field(
"embedding_dimension",
"custom dimension must not exceed 4096",
));
}
}
Ok(())
}
pub fn dimension(&self) -> usize {
self.embedding_dimension.size()
}
}
#[derive(Clone, Debug)]
pub enum EmbeddingProvider {
Builtin {
model_path: Option<PathBuf>,
},
External,
}
impl EmbeddingProvider {
pub fn is_builtin(&self) -> bool {
matches!(self, Self::Builtin { .. })
}
pub fn is_external(&self) -> bool {
matches!(self, Self::External)
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub enum EmbeddingDimension {
#[default]
D384,
D768,
Custom(usize),
}
impl EmbeddingDimension {
#[inline]
pub const fn size(&self) -> usize {
match self {
Self::D384 => 384,
Self::D768 => 768,
Self::Custom(n) => *n,
}
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub enum SyncMode {
#[default]
Normal,
Fast,
Paranoid,
}
impl SyncMode {
pub fn is_paranoid(&self) -> bool {
matches!(self, Self::Paranoid)
}
pub fn is_fast(&self) -> bool {
matches!(self, Self::Fast)
}
}
#[derive(Clone, Debug)]
pub struct HnswConfig {
pub max_nb_connection: usize,
pub ef_construction: usize,
pub ef_search: usize,
pub max_layer: usize,
pub max_elements: usize,
}
impl Default for HnswConfig {
fn default() -> Self {
Self {
max_nb_connection: 16,
ef_construction: 200,
ef_search: 50,
max_layer: 16,
max_elements: 10_000,
}
}
}
#[derive(Clone, Debug)]
pub struct ActivityConfig {
pub stale_threshold: Duration,
}
impl Default for ActivityConfig {
fn default() -> Self {
Self {
stale_threshold: Duration::from_secs(300),
}
}
}
#[derive(Clone, Debug)]
pub struct WatchConfig {
pub in_process: bool,
pub buffer_size: usize,
pub poll_interval_ms: u64,
}
impl Default for WatchConfig {
fn default() -> Self {
Self {
in_process: true,
buffer_size: 1000,
poll_interval_ms: 100,
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
pub struct RecallWeights {
pub similarity: f32,
pub energy: f32,
}
impl RecallWeights {
pub const fn new(similarity: f32, energy: f32) -> Self {
Self { similarity, energy }
}
pub(crate) fn validate(&self, field: &'static str) -> Result<(), ValidationError> {
if !self.similarity.is_finite()
|| !self.energy.is_finite()
|| self.similarity < 0.0
|| self.energy < 0.0
{
return Err(ValidationError::invalid_field(
field,
"weights must be finite and non-negative",
));
}
let sum = self.similarity + self.energy;
if (sum - 1.0).abs() > 0.000_1 {
return Err(ValidationError::invalid_field(
field,
"similarity and energy weights must sum to 1",
));
}
Ok(())
}
}
impl Default for RecallWeights {
fn default() -> Self {
Self {
similarity: 0.7,
energy: 0.3,
}
}
}
#[derive(Clone, Debug, PartialEq)]
pub struct DecayConfig {
pub half_life: Duration,
pub freq_weight: f32,
pub floor: f32,
pub auto_archive_below_floor: bool,
pub default_recall_weights: Option<RecallWeights>,
}
impl Default for DecayConfig {
fn default() -> Self {
Self {
half_life: Duration::from_secs(30 * 24 * 60 * 60),
freq_weight: 0.25,
floor: 0.05,
auto_archive_below_floor: false,
default_recall_weights: None,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_default_config() {
let config = Config::default();
assert!(config.embedding_provider.is_external());
assert_eq!(config.embedding_dimension, EmbeddingDimension::D384);
assert_eq!(config.cache_size_mb, 64);
assert_eq!(config.sync_mode, SyncMode::Normal);
assert!(config.default_collective.is_none());
}
#[test]
fn test_with_builtin_embeddings() {
let config = Config::with_builtin_embeddings();
assert!(config.embedding_provider.is_builtin());
}
#[test]
fn test_with_external_embeddings() {
let config = Config::with_external_embeddings(EmbeddingDimension::Custom(1536));
assert!(config.embedding_provider.is_external());
assert_eq!(config.dimension(), 1536);
}
#[test]
fn test_validate_success() {
let config = Config::default();
assert!(config.validate().is_ok());
}
#[test]
fn test_validate_cache_size_zero() {
let config = Config {
cache_size_mb: 0,
..Default::default()
};
let err = config.validate().unwrap_err();
assert!(
matches!(err, ValidationError::InvalidField { field, .. } if field == "cache_size_mb")
);
}
#[test]
fn test_validate_custom_dimension_zero() {
let config = Config {
embedding_dimension: EmbeddingDimension::Custom(0),
..Default::default()
};
assert!(config.validate().is_err());
}
#[test]
fn test_validate_custom_dimension_too_large() {
let config = Config {
embedding_dimension: EmbeddingDimension::Custom(5000),
..Default::default()
};
assert!(config.validate().is_err());
}
#[test]
fn test_validate_custom_dimension_valid() {
let config = Config {
embedding_dimension: EmbeddingDimension::Custom(1536),
..Default::default()
};
assert!(config.validate().is_ok());
}
#[test]
fn test_embedding_dimension_sizes() {
assert_eq!(EmbeddingDimension::D384.size(), 384);
assert_eq!(EmbeddingDimension::D768.size(), 768);
assert_eq!(EmbeddingDimension::Custom(512).size(), 512);
}
#[test]
fn test_sync_mode_checks() {
assert!(!SyncMode::Normal.is_fast());
assert!(!SyncMode::Normal.is_paranoid());
assert!(SyncMode::Fast.is_fast());
assert!(SyncMode::Paranoid.is_paranoid());
}
#[test]
fn test_hnsw_config_defaults() {
let config = HnswConfig::default();
assert_eq!(config.max_nb_connection, 16);
assert_eq!(config.ef_construction, 200);
assert_eq!(config.ef_search, 50);
assert_eq!(config.max_layer, 16);
assert_eq!(config.max_elements, 10_000);
}
#[test]
fn test_config_includes_hnsw() {
let config = Config::default();
assert_eq!(config.hnsw.max_nb_connection, 16);
}
#[test]
fn test_decay_config_defaults() {
let config = DecayConfig::default();
assert_eq!(config.half_life, Duration::from_secs(30 * 24 * 60 * 60));
assert_eq!(config.freq_weight, 0.25);
assert_eq!(config.floor, 0.05);
assert!(!config.auto_archive_below_floor);
assert!(config.default_recall_weights.is_none());
}
#[test]
fn test_config_includes_decay() {
let config = Config::default();
assert_eq!(
config.decay.half_life,
Duration::from_secs(30 * 24 * 60 * 60)
);
}
#[test]
fn test_validate_decay_zero_half_life() {
let config = Config {
decay: DecayConfig {
half_life: Duration::ZERO,
..Default::default()
},
..Default::default()
};
let err = config.validate().unwrap_err();
assert!(matches!(
err,
ValidationError::InvalidField { field, .. } if field == "decay.half_life"
));
}
#[test]
fn test_validate_decay_negative_freq_weight() {
let config = Config {
decay: DecayConfig {
freq_weight: -0.01,
..Default::default()
},
..Default::default()
};
let err = config.validate().unwrap_err();
assert!(matches!(
err,
ValidationError::InvalidField { field, .. } if field == "decay.freq_weight"
));
}
#[test]
fn test_validate_decay_floor_out_of_range() {
let config = Config {
decay: DecayConfig {
floor: 1.01,
..Default::default()
},
..Default::default()
};
let err = config.validate().unwrap_err();
assert!(matches!(
err,
ValidationError::InvalidField { field, .. } if field == "decay.floor"
));
}
#[test]
fn test_validate_hnsw_zero_max_nb_connection() {
let config = Config {
hnsw: HnswConfig {
max_nb_connection: 0,
..Default::default()
},
..Default::default()
};
let err = config.validate().unwrap_err();
assert!(matches!(
err,
ValidationError::InvalidField { field, .. } if field == "hnsw.max_nb_connection"
));
}
#[test]
fn test_validate_hnsw_zero_ef_construction() {
let config = Config {
hnsw: HnswConfig {
ef_construction: 0,
..Default::default()
},
..Default::default()
};
assert!(config.validate().is_err());
}
#[test]
fn test_validate_hnsw_zero_ef_search() {
let config = Config {
hnsw: HnswConfig {
ef_search: 0,
..Default::default()
},
..Default::default()
};
assert!(config.validate().is_err());
}
#[test]
fn test_embedding_dimension_serialization() {
let dim = EmbeddingDimension::D768;
let bytes = postcard::to_stdvec(&dim).unwrap();
let restored: EmbeddingDimension = postcard::from_bytes(&bytes).unwrap();
assert_eq!(dim, restored);
}
#[test]
fn test_watch_config_defaults() {
let config = WatchConfig::default();
assert!(config.in_process);
assert_eq!(config.buffer_size, 1000);
assert_eq!(config.poll_interval_ms, 100);
}
#[test]
fn test_validate_watch_zero_buffer_size() {
let config = Config {
watch: WatchConfig {
buffer_size: 0,
..Default::default()
},
..Default::default()
};
let err = config.validate().unwrap_err();
assert!(matches!(
err,
ValidationError::InvalidField { field, .. } if field == "watch.buffer_size"
));
}
#[test]
fn test_validate_watch_zero_poll_interval() {
let config = Config {
watch: WatchConfig {
poll_interval_ms: 0,
..Default::default()
},
..Default::default()
};
let err = config.validate().unwrap_err();
assert!(matches!(
err,
ValidationError::InvalidField { field, .. } if field == "watch.poll_interval_ms"
));
}
}