arms/core/
config.rs

1//! # Configuration
2//!
3//! ARMS configuration - define your space.
4//!
5//! Everything is configurable, not hardcoded:
6//! - Dimensionality
7//! - Proximity function
8//! - Merge function
9//! - Tier settings
10//!
11//! "If we say it's a rock now, in 2 years it can never be carved into a wheel."
12
13use super::proximity::{Cosine, Proximity};
14use super::merge::{Mean, Merge};
15use std::sync::Arc;
16
17/// Main ARMS configuration
18///
19/// Defines the dimensional space and default operations.
20#[derive(Clone)]
21pub struct ArmsConfig {
22    /// Dimensionality of the space
23    ///
24    /// Set this to match your model's hidden size.
25    /// Examples: 768 (BERT), 1024 (GPT-2 medium), 4096 (large models)
26    pub dimensionality: usize,
27
28    /// Proximity function for similarity calculations
29    pub proximity: Arc<dyn Proximity>,
30
31    /// Merge function for hierarchical composition
32    pub merge: Arc<dyn Merge>,
33
34    /// Whether to normalize points on insertion
35    pub normalize_on_insert: bool,
36
37    /// Tier configuration
38    pub tiers: TierConfig,
39}
40
41impl ArmsConfig {
42    /// Create a new configuration with specified dimensionality
43    ///
44    /// Uses default proximity (Cosine) and merge (Mean) functions.
45    pub fn new(dimensionality: usize) -> Self {
46        Self {
47            dimensionality,
48            proximity: Arc::new(Cosine),
49            merge: Arc::new(Mean),
50            normalize_on_insert: true,
51            tiers: TierConfig::default(),
52        }
53    }
54
55    /// Set a custom proximity function
56    pub fn with_proximity<P: Proximity + 'static>(mut self, proximity: P) -> Self {
57        self.proximity = Arc::new(proximity);
58        self
59    }
60
61    /// Set a custom merge function
62    pub fn with_merge<M: Merge + 'static>(mut self, merge: M) -> Self {
63        self.merge = Arc::new(merge);
64        self
65    }
66
67    /// Set normalization behavior
68    pub fn with_normalize(mut self, normalize: bool) -> Self {
69        self.normalize_on_insert = normalize;
70        self
71    }
72
73    /// Set tier configuration
74    pub fn with_tiers(mut self, tiers: TierConfig) -> Self {
75        self.tiers = tiers;
76        self
77    }
78}
79
80impl Default for ArmsConfig {
81    /// Default configuration: 768 dimensions, cosine proximity, mean merge
82    fn default() -> Self {
83        Self::new(768)
84    }
85}
86
87/// Tier configuration for storage management
88#[derive(Clone, Debug)]
89pub struct TierConfig {
90    /// Hot tier (RAM) capacity in bytes
91    pub hot_capacity: usize,
92
93    /// Warm tier (NVMe) capacity in bytes
94    pub warm_capacity: usize,
95
96    /// Number of accesses before promoting to hotter tier
97    pub promote_after_accesses: u32,
98
99    /// Milliseconds since last access before evicting to colder tier
100    pub evict_after_ms: u64,
101}
102
103impl TierConfig {
104    /// Create a new tier configuration
105    pub fn new(hot_capacity: usize, warm_capacity: usize) -> Self {
106        Self {
107            hot_capacity,
108            warm_capacity,
109            promote_after_accesses: 3,
110            evict_after_ms: 3600 * 1000, // 1 hour
111        }
112    }
113
114    /// Tiny config for testing
115    pub fn tiny() -> Self {
116        Self {
117            hot_capacity: 1024 * 1024,        // 1 MB
118            warm_capacity: 10 * 1024 * 1024,  // 10 MB
119            promote_after_accesses: 2,
120            evict_after_ms: 60 * 1000, // 1 minute
121        }
122    }
123}
124
125impl Default for TierConfig {
126    fn default() -> Self {
127        Self {
128            hot_capacity: 1024 * 1024 * 1024,       // 1 GB
129            warm_capacity: 100 * 1024 * 1024 * 1024, // 100 GB
130            promote_after_accesses: 3,
131            evict_after_ms: 3600 * 1000, // 1 hour
132        }
133    }
134}
135
136#[cfg(test)]
137mod tests {
138    use super::*;
139    use crate::core::proximity::Euclidean;
140    use crate::core::merge::MaxPool;
141
142    #[test]
143    fn test_default_config() {
144        let config = ArmsConfig::default();
145        assert_eq!(config.dimensionality, 768);
146        assert!(config.normalize_on_insert);
147        assert_eq!(config.proximity.name(), "cosine");
148        assert_eq!(config.merge.name(), "mean");
149    }
150
151    #[test]
152    fn test_custom_config() {
153        let config = ArmsConfig::new(4096)
154            .with_proximity(Euclidean)
155            .with_merge(MaxPool)
156            .with_normalize(false);
157
158        assert_eq!(config.dimensionality, 4096);
159        assert!(!config.normalize_on_insert);
160        assert_eq!(config.proximity.name(), "euclidean");
161        assert_eq!(config.merge.name(), "max_pool");
162    }
163
164    #[test]
165    fn test_tier_config() {
166        let tiers = TierConfig::new(1024, 2048);
167        assert_eq!(tiers.hot_capacity, 1024);
168        assert_eq!(tiers.warm_capacity, 2048);
169    }
170
171    #[test]
172    fn test_tier_tiny() {
173        let tiers = TierConfig::tiny();
174        assert_eq!(tiers.hot_capacity, 1024 * 1024);
175        assert_eq!(tiers.evict_after_ms, 60 * 1000);
176    }
177}
arms/core/config.rs

arms/core/
config.rs
