nt_memory/
lib.rs

1//! Neural Trader Memory Systems
2//!
3//! Three-tier memory hierarchy for high-performance agent coordination:
4//!
5//! - **L1 Cache**: Hot cache with DashMap (<1μs lookup)
6//! - **L2 Vector DB**: AgentDB for semantic search (<1ms p95)
7//! - **L3 Cold Storage**: Sled embedded database (>10ms)
8//!
9//! ## Features
10//!
11//! - Vector embeddings and semantic search
12//! - ReasoningBank trajectory tracking
13//! - Cross-agent coordination with pub/sub
14//! - Distributed locks and consensus
15//! - Memory distillation and compression
16//! - Session persistence
17//!
18//! ## Performance Targets
19//!
20//! - L1 cache: <1μs lookup
21//! - Vector search: <1ms (p95)
22//! - Position lookup: <100ns (p99)
23//! - Memory footprint: <1GB for 1M observations
24//! - Cross-agent latency: <5ms
25
26pub mod cache;
27pub mod agentdb;
28pub mod reasoningbank;
29pub mod coordination;
30
31// Re-exports for convenient access
32pub use cache::{HotCache, CacheConfig, CacheEntry};
33pub use agentdb::{VectorStore, EmbeddingProvider};
34pub use reasoningbank::{
35    TrajectoryTracker,
36    VerdictJudge,
37    MemoryDistiller,
38    Trajectory,
39    Verdict,
40};
41pub use coordination::{
42    PubSubBroker,
43    DistributedLock,
44    ConsensusEngine,
45    Namespace,
46};
47
48use thiserror::Error;
49
50/// Memory system errors
51#[derive(Debug, Error)]
52pub enum MemoryError {
53    #[error("Cache error: {0}")]
54    Cache(String),
55
56    #[error("Vector database error: {0}")]
57    VectorDB(String),
58
59    #[error("Storage error: {0}")]
60    Storage(String),
61
62    #[error("Coordination error: {0}")]
63    Coordination(String),
64
65    #[error("Serialization error: {0}")]
66    Serialization(String),
67
68    #[error("Lock acquisition timeout")]
69    LockTimeout,
70
71    #[error("Consensus not reached")]
72    ConsensusFailure,
73
74    #[error("Invalid namespace: {0}")]
75    InvalidNamespace(String),
76}
77
78pub type Result<T> = std::result::Result<T, MemoryError>;
79
80/// Memory system configuration
81#[derive(Debug, Clone)]
82pub struct MemoryConfig {
83    /// L1 cache configuration
84    pub cache_config: CacheConfig,
85
86    /// AgentDB base URL
87    pub agentdb_url: String,
88
89    /// Cold storage path
90    pub storage_path: String,
91
92    /// Enable compression for distillation
93    pub enable_compression: bool,
94
95    /// Maximum memory footprint (bytes)
96    pub max_memory_bytes: usize,
97}
98
99impl Default for MemoryConfig {
100    fn default() -> Self {
101        Self {
102            cache_config: CacheConfig::default(),
103            agentdb_url: "http://localhost:3000".to_string(),
104            storage_path: "./data/memory".to_string(),
105            enable_compression: true,
106            max_memory_bytes: 1_073_741_824, // 1GB
107        }
108    }
109}
110
111/// Unified memory system interface
112pub struct MemorySystem {
113    /// L1: Hot cache
114    cache: HotCache,
115
116    /// L2: Vector database
117    vector_store: VectorStore,
118
119    /// L3: Cold storage
120    cold_storage: sled::Db,
121
122    /// ReasoningBank components
123    trajectory_tracker: TrajectoryTracker,
124    verdict_judge: VerdictJudge,
125    distiller: MemoryDistiller,
126
127    /// Coordination components
128    pubsub: PubSubBroker,
129    locks: DistributedLock,
130    consensus: ConsensusEngine,
131
132    /// Configuration
133    config: MemoryConfig,
134}
135
136impl MemorySystem {
137    /// Create new memory system
138    pub async fn new(config: MemoryConfig) -> Result<Self> {
139        // Initialize L1 cache
140        let cache = HotCache::new(config.cache_config.clone());
141
142        // Initialize L2 vector store
143        let vector_store = VectorStore::new(&config.agentdb_url)
144            .await
145            .map_err(|e| MemoryError::VectorDB(e.to_string()))?;
146
147        // Initialize L3 cold storage
148        let cold_storage = sled::open(&config.storage_path)
149            .map_err(|e| MemoryError::Storage(e.to_string()))?;
150
151        // Initialize ReasoningBank
152        let trajectory_tracker = TrajectoryTracker::new();
153        let verdict_judge = VerdictJudge::new();
154        let distiller = MemoryDistiller::new(config.enable_compression);
155
156        // Initialize coordination
157        let pubsub = PubSubBroker::new();
158        let locks = DistributedLock::new();
159        let consensus = ConsensusEngine::new();
160
161        Ok(Self {
162            cache,
163            vector_store,
164            cold_storage,
165            trajectory_tracker,
166            verdict_judge,
167            distiller,
168            pubsub,
169            locks,
170            consensus,
171            config,
172        })
173    }
174
175    /// Get from memory (tries L1 -> L2 -> L3)
176    pub async fn get(&self, namespace: &str, key: &str) -> Result<Option<Vec<u8>>> {
177        let full_key = format!("{}/{}", namespace, key);
178
179        // Try L1 cache
180        if let Some(entry) = self.cache.get(&full_key) {
181            tracing::debug!("L1 cache hit: {}", full_key);
182            return Ok(Some(entry.data));
183        }
184
185        // Try L3 cold storage
186        if let Some(data) = self.cold_storage
187            .get(full_key.as_bytes())
188            .map_err(|e| MemoryError::Storage(e.to_string()))?
189        {
190            tracing::debug!("L3 storage hit: {}", full_key);
191
192            // Promote to L1
193            self.cache.insert(&full_key, data.to_vec());
194
195            return Ok(Some(data.to_vec()));
196        }
197
198        tracing::debug!("Cache miss: {}", full_key);
199        Ok(None)
200    }
201
202    /// Store in memory (writes to all tiers)
203    pub async fn put(&self, namespace: &str, key: &str, value: Vec<u8>) -> Result<()> {
204        let full_key = format!("{}/{}", namespace, key);
205
206        // Write to L1
207        self.cache.insert(&full_key, value.clone());
208
209        // Write to L3 (async)
210        self.cold_storage
211            .insert(full_key.as_bytes(), value.as_slice())
212            .map_err(|e| MemoryError::Storage(e.to_string()))?;
213
214        tracing::debug!("Stored: {}", full_key);
215        Ok(())
216    }
217
218    /// Semantic search using L2 vector store
219    pub async fn search_similar(
220        &self,
221        namespace: &str,
222        query_embedding: Vec<f32>,
223        top_k: usize,
224    ) -> Result<Vec<(String, f32)>> {
225        self.vector_store
226            .search(namespace, query_embedding, top_k)
227            .await
228            .map_err(|e| MemoryError::VectorDB(e.to_string()))
229    }
230
231    /// Track agent trajectory
232    pub async fn track_trajectory(&self, trajectory: Trajectory) -> Result<()> {
233        self.trajectory_tracker
234            .track(trajectory)
235            .await
236            .map_err(|e| MemoryError::Storage(e.to_string()))
237    }
238
239    /// Subscribe to agent messages
240    pub async fn subscribe(&self, topic: &str) -> Result<tokio::sync::mpsc::Receiver<Vec<u8>>> {
241        self.pubsub
242            .subscribe(topic)
243            .await
244            .map_err(|e| MemoryError::Coordination(e.to_string()))
245    }
246
247    /// Publish message to agents
248    pub async fn publish(&self, topic: &str, message: Vec<u8>) -> Result<()> {
249        self.pubsub
250            .publish(topic, message)
251            .await
252            .map_err(|e| MemoryError::Coordination(e.to_string()))
253    }
254
255    /// Acquire distributed lock
256    pub async fn acquire_lock(
257        &self,
258        resource: &str,
259        timeout: std::time::Duration,
260    ) -> Result<String> {
261        self.locks
262            .acquire(resource, timeout)
263            .await
264            .map_err(|e| MemoryError::Coordination(e.to_string()))
265    }
266
267    /// Release distributed lock
268    pub async fn release_lock(&self, token: &str) -> Result<()> {
269        self.locks
270            .release(token)
271            .await
272            .map_err(|e| MemoryError::Coordination(e.to_string()))
273    }
274
275    /// Get memory statistics
276    pub fn stats(&self) -> MemoryStats {
277        MemoryStats {
278            l1_entries: self.cache.len(),
279            l1_hit_rate: self.cache.hit_rate(),
280            l3_size_bytes: self.cold_storage.size_on_disk().unwrap_or(0),
281            total_trajectories: self.trajectory_tracker.count(),
282        }
283    }
284}
285
286/// Memory system statistics
287#[derive(Debug, Clone)]
288pub struct MemoryStats {
289    pub l1_entries: usize,
290    pub l1_hit_rate: f64,
291    pub l3_size_bytes: u64,
292    pub total_trajectories: usize,
293}
294
295#[cfg(test)]
296mod tests {
297    use super::*;
298
299    #[tokio::test]
300    async fn test_memory_system_creation() {
301        let mem_config = MemoryConfig {
302            storage_path: tempfile::tempdir().unwrap().path().to_str().unwrap().to_string(),
303            ..Default::default()
304        };
305
306        let memory = MemorySystem::new(mem_config).await;
307        assert!(memory.is_ok());
308    }
309}