mod config;
mod long_term;
mod short_term;
mod surprise_gate;
pub use config::IneruConfig;
pub use long_term::{LongTermMemory, MemoryBank};
pub use short_term::ShortTermMemory;
pub use surprise_gate::SurpriseGate;
use crate::error::AiResult;
use crate::types::{
AiTransaction, AnomalyResult, MemoryMatch, MemorySource, Pattern, ProcessResult,
};
use parking_lot::RwLock;
use std::sync::Arc;
use tracing::{debug, trace};
pub struct IneruMemory {
short_term: Arc<RwLock<ShortTermMemory>>,
long_term: Arc<RwLock<LongTermMemory>>,
surprise_gate: SurpriseGate,
config: IneruConfig,
}
impl IneruMemory {
pub fn new(config: IneruConfig) -> Self {
let short_term = ShortTermMemory::new(config.window_size);
let long_term = LongTermMemory::new(config.memory_capacity, config.embedding_dim);
let surprise_gate = SurpriseGate::new(config.surprise_threshold);
Self {
short_term: Arc::new(RwLock::new(short_term)),
long_term: Arc::new(RwLock::new(long_term)),
surprise_gate,
config,
}
}
pub fn process(&mut self, tx: &AiTransaction) -> AiResult<ProcessResult> {
let pattern = tx.to_pattern();
trace!(hash = ?tx.hash, "Processing transaction through Ineru memory");
{
let mut stm = self.short_term.write();
stm.add(pattern.clone());
}
let relevance = {
let stm = self.short_term.read();
stm.attention_score(&pattern)
};
let surprise = {
let ltm = self.long_term.read();
self.surprise_gate.compute_surprise(&pattern, <m)
};
let stored_long_term = surprise > self.config.surprise_threshold;
if stored_long_term {
debug!(
surprise = surprise,
threshold = self.config.surprise_threshold,
"Storing pattern in long-term memory"
);
let mut ltm = self.long_term.write();
ltm.update(pattern.clone())?;
self.surprise_gate.observe(&pattern);
}
let anomaly = if self.config.anomaly_detection {
Some(self.detect_anomaly(tx)?)
} else {
None
};
Ok(ProcessResult {
relevance,
surprise,
stored_long_term,
anomaly,
})
}
pub fn query(&self, pattern: &Pattern, limit: usize) -> Vec<MemoryMatch> {
let mut matches = Vec::new();
{
let stm = self.short_term.read();
let stm_matches = stm.search(&pattern.embedding, limit);
for (p, score) in stm_matches {
matches.push(MemoryMatch {
pattern: p,
similarity: score,
source: MemorySource::ShortTerm,
});
}
}
{
let ltm = self.long_term.read();
let ltm_matches = ltm.retrieve(&pattern.embedding, limit);
for (p, score) in ltm_matches {
matches.push(MemoryMatch {
pattern: p,
similarity: score,
source: MemorySource::LongTerm,
});
}
}
matches.sort_by(|a, b| b.similarity.partial_cmp(&a.similarity).unwrap());
matches.truncate(limit);
matches
}
pub fn detect_anomaly(&self, tx: &AiTransaction) -> AiResult<AnomalyResult> {
let pattern = tx.to_pattern();
let similarity_short = {
let stm = self.short_term.read();
stm.max_similarity(&pattern.embedding)
};
let similarity_long = {
let ltm = self.long_term.read();
ltm.max_similarity(&pattern.embedding)
};
let max_similarity = similarity_short.max(similarity_long);
let is_anomaly = max_similarity < self.config.anomaly_threshold;
let confidence = 1.0 - max_similarity;
let reason = if is_anomaly {
format!(
"Transaction pattern differs from known patterns (similarity: {:.2})",
max_similarity
)
} else {
"Transaction pattern matches known patterns".to_string()
};
Ok(AnomalyResult {
is_anomaly,
confidence,
reason,
})
}
pub fn stats(&self) -> MemoryStats {
let stm = self.short_term.read();
let ltm = self.long_term.read();
MemoryStats {
short_term_size: stm.len(),
short_term_capacity: self.config.window_size,
long_term_size: ltm.len(),
long_term_capacity: self.config.memory_capacity,
total_patterns: stm.len() + ltm.len(),
}
}
pub fn clear(&self) {
self.short_term.write().clear();
self.long_term.write().clear();
}
pub fn predict_outcome(&self, tx: &AiTransaction) -> PredictionResult {
let pattern = tx.to_pattern();
let matches = self.query(&pattern, 10);
if matches.is_empty() {
return PredictionResult {
confidence: 0.0,
predicted_valid: true, similar_count: 0,
};
}
let weighted_sum: f32 = matches.iter().map(|m| m.similarity).sum();
PredictionResult {
confidence: (weighted_sum / matches.len() as f32).min(1.0),
predicted_valid: true,
similar_count: matches.len(),
}
}
}
#[derive(Debug, Clone)]
pub struct MemoryStats {
pub short_term_size: usize,
pub short_term_capacity: usize,
pub long_term_size: usize,
pub long_term_capacity: usize,
pub total_patterns: usize,
}
#[derive(Debug, Clone)]
pub struct PredictionResult {
pub confidence: f32,
pub predicted_valid: bool,
pub similar_count: usize,
}
#[cfg(test)]
mod tests {
use super::*;
fn make_test_tx(id: u8) -> AiTransaction {
AiTransaction {
hash: [id; 32],
timestamp: 1702656000000 + (id as u64 * 1000),
agent: [1u8; 32],
entry_type: "test".to_string(),
data: vec![id; 10],
size: 10,
}
}
#[test]
fn test_ineru_memory_basic() {
let config = IneruConfig::default();
let mut memory = IneruMemory::new(config);
let tx = make_test_tx(1);
let result = memory.process(&tx).unwrap();
assert!(result.relevance >= 0.0);
assert!(result.surprise >= 0.0);
}
#[test]
fn test_ineru_memory_query() {
let config = IneruConfig::default();
let mut memory = IneruMemory::new(config);
for i in 0..10 {
let tx = make_test_tx(i);
memory.process(&tx).unwrap();
}
let query_tx = make_test_tx(5);
let pattern = query_tx.to_pattern();
let matches = memory.query(&pattern, 5);
assert!(!matches.is_empty());
}
#[test]
fn test_anomaly_detection() {
let config = IneruConfig {
anomaly_detection: true,
anomaly_threshold: 0.5,
..IneruConfig::default()
};
let mut memory = IneruMemory::new(config);
for i in 0..20 {
let tx = make_test_tx(i % 3); memory.process(&tx).unwrap();
}
let known_tx = make_test_tx(1);
let _result = memory.detect_anomaly(&known_tx).unwrap();
let anomaly_tx = AiTransaction {
hash: [255u8; 32],
timestamp: 1702656000000,
agent: [255u8; 32],
entry_type: "unknown_type".to_string(),
data: vec![255; 100],
size: 100,
};
let result = memory.detect_anomaly(&anomaly_tx).unwrap();
assert!(result.confidence >= 0.0);
}
#[test]
fn test_memory_stats() {
let config = IneruConfig::default();
let mut memory = IneruMemory::new(config.clone());
let stats = memory.stats();
assert_eq!(stats.short_term_size, 0);
assert_eq!(stats.short_term_capacity, config.window_size);
for i in 0..5 {
let tx = make_test_tx(i);
memory.process(&tx).unwrap();
}
let stats = memory.stats();
assert_eq!(stats.short_term_size, 5);
}
}