use heapless::Vec as HVec;
use heapless::String as HString;
use super::{MicroHNSW, HNSWConfig, SearchResult, MicroVector, DistanceMetric};
pub const MAX_MEMORIES: usize = 128;
pub const MAX_TEXT_LEN: usize = 64;
pub const MEMORY_DIM: usize = 32;
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum MemoryType {
Preference,
Fact,
Event,
Procedure,
Entity,
Emotion,
Context,
State,
}
impl MemoryType {
pub fn priority(&self) -> i32 {
match self {
Self::State => 100, Self::Context => 90, Self::Preference => 80, Self::Emotion => 70, Self::Procedure => 60, Self::Fact => 50, Self::Event => 40, Self::Entity => 30, }
}
}
#[derive(Debug, Clone)]
pub struct Memory {
pub id: u32,
pub memory_type: MemoryType,
pub timestamp: u32,
pub text: HString<MAX_TEXT_LEN>,
pub importance: u8,
pub access_count: u16,
pub embedding: HVec<i8, MEMORY_DIM>,
}
impl Memory {
pub fn new(
id: u32,
memory_type: MemoryType,
text: &str,
embedding: &[i8],
timestamp: u32,
) -> Option<Self> {
let mut text_str = HString::new();
for c in text.chars().take(MAX_TEXT_LEN) {
text_str.push(c).ok()?;
}
let mut embed_vec = HVec::new();
for &v in embedding.iter().take(MEMORY_DIM) {
embed_vec.push(v).ok()?;
}
Some(Self {
id,
memory_type,
timestamp,
text: text_str,
importance: 50,
access_count: 0,
embedding: embed_vec,
})
}
pub fn relevance_score(&self, distance: i32, current_time: u32) -> i32 {
let type_weight = self.memory_type.priority();
let importance_weight = self.importance as i32;
let age_seconds = current_time.saturating_sub(self.timestamp);
let recency = 100 - (age_seconds / 3600).min(100) as i32;
let frequency = (self.access_count as i32).min(50);
let distance_score = 1000 - distance.min(1000);
(distance_score * 3 + type_weight * 2 + importance_weight + recency + frequency) / 7
}
}
pub struct SemanticMemory {
index: MicroHNSW<MEMORY_DIM, MAX_MEMORIES>,
memories: HVec<Memory, MAX_MEMORIES>,
next_id: u32,
current_time: u32,
}
impl SemanticMemory {
pub fn new() -> Self {
let config = HNSWConfig {
m: 4,
m_max0: 8,
ef_construction: 16,
ef_search: 8,
metric: DistanceMetric::Euclidean,
binary_mode: false,
};
Self {
index: MicroHNSW::new(config),
memories: HVec::new(),
next_id: 0,
current_time: 0,
}
}
pub fn set_time(&mut self, time: u32) {
self.current_time = time;
}
pub fn len(&self) -> usize {
self.memories.len()
}
pub fn is_empty(&self) -> bool {
self.memories.is_empty()
}
pub fn memory_bytes(&self) -> usize {
self.index.memory_bytes() + self.memories.len() * core::mem::size_of::<Memory>()
}
pub fn remember(
&mut self,
memory_type: MemoryType,
text: &str,
embedding: &[i8],
) -> Result<u32, &'static str> {
if self.memories.len() >= MAX_MEMORIES {
self.evict_least_important()?;
}
let id = self.next_id;
self.next_id += 1;
let memory = Memory::new(id, memory_type, text, embedding, self.current_time)
.ok_or("Failed to create memory")?;
let vec = MicroVector {
data: memory.embedding.clone(),
id,
};
self.index.insert(&vec)?;
self.memories.push(memory).map_err(|_| "Memory full")?;
Ok(id)
}
pub fn recall(&mut self, query_embedding: &[i8], k: usize) -> HVec<(Memory, i32), 16> {
let mut results = HVec::new();
let search_results = self.index.search(query_embedding, k * 2);
for result in search_results.iter() {
if let Some(memory) = self.find_memory_by_id(result.id) {
let score = memory.relevance_score(result.distance, self.current_time);
let _ = results.push((memory.clone(), score));
}
}
results.sort_by(|a, b| b.1.cmp(&a.1));
for (mem, _) in results.iter() {
self.increment_access(mem.id);
}
while results.len() > k {
results.pop();
}
results
}
pub fn recall_by_type(
&mut self,
query_embedding: &[i8],
memory_type: MemoryType,
k: usize,
) -> HVec<Memory, 16> {
let all_results = self.recall(query_embedding, k * 3);
let mut filtered = HVec::new();
for (memory, _) in all_results {
if memory.memory_type == memory_type && filtered.len() < k {
let _ = filtered.push(memory);
}
}
filtered
}
pub fn recent(&self, k: usize) -> HVec<&Memory, 16> {
let mut sorted: HVec<&Memory, MAX_MEMORIES> = self.memories.iter().collect();
sorted.sort_by(|a, b| b.timestamp.cmp(&a.timestamp));
let mut result = HVec::new();
for mem in sorted.iter().take(k) {
let _ = result.push(*mem);
}
result
}
pub fn forget(&mut self, id: u32) -> bool {
if let Some(pos) = self.memories.iter().position(|m| m.id == id) {
self.memories.swap_remove(pos);
true
} else {
false
}
}
fn find_memory_by_id(&self, id: u32) -> Option<&Memory> {
self.memories.iter().find(|m| m.id == id)
}
fn increment_access(&mut self, id: u32) {
if let Some(memory) = self.memories.iter_mut().find(|m| m.id == id) {
memory.access_count = memory.access_count.saturating_add(1);
}
}
fn evict_least_important(&mut self) -> Result<(), &'static str> {
if self.memories.is_empty() {
return Ok(());
}
let mut min_score = i32::MAX;
let mut min_idx = 0;
for (i, memory) in self.memories.iter().enumerate() {
let score = memory.relevance_score(0, self.current_time);
if score < min_score {
min_score = score;
min_idx = i;
}
}
self.memories.swap_remove(min_idx);
Ok(())
}
}
impl Default for SemanticMemory {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_memory_creation() {
let embedding = [10i8; MEMORY_DIM];
let memory = Memory::new(1, MemoryType::Preference, "I like coffee", &embedding, 1000);
assert!(memory.is_some());
let m = memory.unwrap();
assert_eq!(m.id, 1);
assert_eq!(m.memory_type, MemoryType::Preference);
}
#[test]
fn test_semantic_memory() {
let mut sm = SemanticMemory::new();
sm.set_time(1000);
let embed1 = [10i8; MEMORY_DIM];
let embed2 = [20i8; MEMORY_DIM];
sm.remember(MemoryType::Preference, "I like tea", &embed1).unwrap();
sm.remember(MemoryType::Fact, "Water is wet", &embed2).unwrap();
assert_eq!(sm.len(), 2);
let query = [11i8; MEMORY_DIM];
let results = sm.recall(&query, 1);
assert!(!results.is_empty());
}
}