use serde::{Deserialize, Serialize};
use crate::error::Result;
use crate::model::memory::MemoryRecord;
use crate::query::MnemoEngine;
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
pub struct MaturityWeights {
pub recency: f32,
pub hit_success: f32,
pub edge_degree: f32,
pub redundancy: f32,
}
impl MaturityWeights {
pub const fn balanced() -> Self {
Self {
recency: 0.25,
hit_success: 0.30,
edge_degree: 0.20,
redundancy: 0.25,
}
}
fn clamped(self) -> Self {
Self {
recency: self.recency.clamp(0.0, 1.0),
hit_success: self.hit_success.clamp(0.0, 1.0),
edge_degree: self.edge_degree.clamp(0.0, 1.0),
redundancy: self.redundancy.clamp(0.0, 1.0),
}
}
fn sum(self) -> f32 {
self.recency + self.hit_success + self.edge_degree + self.redundancy
}
}
impl Default for MaturityWeights {
fn default() -> Self {
Self::balanced()
}
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
pub struct MaturityBreakdown {
pub recency: f32,
pub hit_success: f32,
pub edge_degree: f32,
pub redundancy: f32,
pub combined: f32,
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
pub struct MaturitySaturation {
pub recency_half_life_hours: f32,
pub hit_saturation: f32,
pub degree_saturation: f32,
}
impl Default for MaturitySaturation {
fn default() -> Self {
Self {
recency_half_life_hours: 72.0,
hit_saturation: 8.0,
degree_saturation: 6.0,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct MaturityPolicy {
pub weights: MaturityWeights,
pub saturation: MaturitySaturation,
pub threshold: f32,
pub min_cluster_size_floor: usize,
pub trigger_on_forget: bool,
pub trigger_on_checkpoint: bool,
}
impl MaturityPolicy {
pub fn balanced() -> Self {
Self {
weights: MaturityWeights::balanced(),
saturation: MaturitySaturation::default(),
threshold: 0.55,
min_cluster_size_floor: 2,
trigger_on_forget: true,
trigger_on_checkpoint: false,
}
}
}
impl Default for MaturityPolicy {
fn default() -> Self {
Self::balanced()
}
}
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum ConsolidationPolicy {
#[default]
FixedSize,
MaturityDriven(MaturityPolicy),
}
pub async fn compute_cluster_maturity(
engine: &MnemoEngine,
cluster: &[&MemoryRecord],
weights: MaturityWeights,
saturation: MaturitySaturation,
) -> Result<Option<MaturityBreakdown>> {
if cluster.is_empty() {
return Ok(None);
}
let weights = weights.clamped();
let recency = recency_component(cluster, saturation.recency_half_life_hours);
let hit_success = hit_success_component(cluster, saturation.hit_saturation);
let edge_degree = edge_degree_component(engine, cluster, saturation.degree_saturation).await?;
let redundancy = redundancy_component(cluster);
let combined = combined_score(
weights,
MaturityBreakdown {
recency,
hit_success,
edge_degree,
redundancy,
combined: 0.0,
},
);
Ok(Some(MaturityBreakdown {
recency,
hit_success,
edge_degree,
redundancy,
combined,
}))
}
fn combined_score(weights: MaturityWeights, b: MaturityBreakdown) -> f32 {
let total = weights.sum();
if total <= f32::EPSILON {
return 0.0;
}
let mixed = weights.recency * b.recency
+ weights.hit_success * b.hit_success
+ weights.edge_degree * b.edge_degree
+ weights.redundancy * b.redundancy;
(mixed / total).clamp(0.0, 1.0)
}
fn recency_component(cluster: &[&MemoryRecord], half_life_hours: f32) -> f32 {
if cluster.is_empty() {
return 0.0;
}
let half_life = half_life_hours.max(f32::EPSILON);
let lambda = std::f32::consts::LN_2 / half_life;
let now = chrono::Utc::now();
let mut sum = 0.0_f32;
let mut n = 0_u32;
for r in cluster {
let last = r.last_accessed_at.as_deref().unwrap_or(&r.created_at);
if let Ok(dt) = chrono::DateTime::parse_from_rfc3339(last) {
let hours = (now - dt.with_timezone(&chrono::Utc)).num_seconds().max(0) as f32 / 3600.0;
sum += (-lambda * hours).exp();
n += 1;
}
}
if n == 0 { 0.0 } else { sum / n as f32 }
}
fn hit_success_component(cluster: &[&MemoryRecord], hit_saturation: f32) -> f32 {
if cluster.is_empty() {
return 0.0;
}
let sat = hit_saturation.max(1.0);
let denom = (1.0 + sat).ln().max(f32::EPSILON);
let mut sum = 0.0_f32;
for r in cluster {
sum += (1.0 + r.access_count as f32).ln() / denom;
}
(sum / cluster.len() as f32).clamp(0.0, 1.0)
}
async fn edge_degree_component(
engine: &MnemoEngine,
cluster: &[&MemoryRecord],
degree_saturation: f32,
) -> Result<f32> {
if cluster.is_empty() {
return Ok(0.0);
}
let sat = degree_saturation.max(1.0);
let mut total_degree = 0_usize;
for r in cluster {
let outgoing = engine.storage.get_relations_from(r.id).await?.len();
let incoming = engine.storage.get_relations_to(r.id).await?.len();
total_degree += outgoing + incoming;
}
let mean = total_degree as f32 / cluster.len() as f32;
Ok((mean / sat).clamp(0.0, 1.0))
}
fn redundancy_component(cluster: &[&MemoryRecord]) -> f32 {
let with_emb: Vec<&Vec<f32>> = cluster
.iter()
.filter_map(|r| r.embedding.as_ref())
.collect();
if with_emb.len() < 2 {
return 0.5;
}
let mut sum = 0.0_f32;
let mut pairs = 0_u32;
for i in 0..with_emb.len() {
for j in (i + 1)..with_emb.len() {
if with_emb[i].len() != with_emb[j].len() || with_emb[i].is_empty() {
continue;
}
sum += cosine_similarity(with_emb[i], with_emb[j]);
pairs += 1;
}
}
if pairs == 0 {
0.5
} else {
(sum / pairs as f32).clamp(0.0, 1.0)
}
}
fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
let mut dot = 0.0_f32;
let mut na = 0.0_f32;
let mut nb = 0.0_f32;
for i in 0..a.len() {
dot += a[i] * b[i];
na += a[i] * a[i];
nb += b[i] * b[i];
}
let denom = (na.sqrt() * nb.sqrt()).max(f32::EPSILON);
dot / denom
}
#[cfg(test)]
mod tests {
use super::*;
use crate::embedding::NoopEmbedding;
use crate::index::usearch::UsearchIndex;
use crate::model::memory::{ConsolidationState, MemoryType, Scope, SourceType};
use crate::storage::duckdb::DuckDbStorage;
use std::sync::Arc;
use uuid::Uuid;
fn record(now: chrono::DateTime<chrono::Utc>, hours_ago: i64, access: u64) -> MemoryRecord {
let created = (now - chrono::Duration::hours(hours_ago)).to_rfc3339();
MemoryRecord {
id: Uuid::now_v7(),
agent_id: "a".to_string(),
content: "c".to_string(),
memory_type: MemoryType::Episodic,
scope: Scope::Private,
importance: 0.5,
tags: vec![],
metadata: serde_json::json!({}),
embedding: None,
content_hash: vec![],
prev_hash: None,
source_type: SourceType::Agent,
source_id: None,
consolidation_state: ConsolidationState::Raw,
access_count: access,
org_id: None,
thread_id: None,
created_at: created.clone(),
updated_at: created,
last_accessed_at: None,
expires_at: None,
deleted_at: None,
decay_rate: None,
created_by: None,
version: 1,
prev_version_id: None,
quarantined: false,
quarantine_reason: None,
decay_function: None,
}
}
#[test]
fn recency_decays_with_age() {
let now = chrono::Utc::now();
let fresh = record(now, 0, 0);
let stale = record(now, 200, 0);
let fresh_score = recency_component(&[&fresh], 72.0);
let stale_score = recency_component(&[&stale], 72.0);
assert!(
fresh_score > stale_score,
"fresh {fresh_score} > stale {stale_score}"
);
let half = record(now, 72, 0);
let half_score = recency_component(&[&half], 72.0);
assert!(
(half_score - 0.5).abs() < 0.05,
"half-life mapping: {half_score} ≈ 0.5"
);
}
#[test]
fn hit_success_saturates() {
let now = chrono::Utc::now();
let none = record(now, 0, 0);
let some = record(now, 0, 4);
let many = record(now, 0, 64);
let s0 = hit_success_component(&[&none], 8.0);
let s1 = hit_success_component(&[&some], 8.0);
let s2 = hit_success_component(&[&many], 8.0);
assert!(s0 < s1 && s1 <= s2);
assert!(s2 <= 1.0);
assert_eq!(s0, 0.0);
}
#[test]
fn redundancy_handles_short_input() {
let now = chrono::Utc::now();
let r = record(now, 0, 0);
let s = redundancy_component(&[&r]);
assert_eq!(s, 0.5);
}
#[test]
fn redundancy_detects_identical_embeddings() {
let now = chrono::Utc::now();
let mut a = record(now, 0, 0);
let mut b = record(now, 0, 0);
a.embedding = Some(vec![1.0, 0.0, 0.0]);
b.embedding = Some(vec![1.0, 0.0, 0.0]);
let s = redundancy_component(&[&a, &b]);
assert!((s - 1.0).abs() < 1e-5, "identical → 1.0, got {s}");
}
#[test]
fn redundancy_orthogonal_is_zero() {
let now = chrono::Utc::now();
let mut a = record(now, 0, 0);
let mut b = record(now, 0, 0);
a.embedding = Some(vec![1.0, 0.0]);
b.embedding = Some(vec![0.0, 1.0]);
let s = redundancy_component(&[&a, &b]);
assert!(s.abs() < 1e-5, "orthogonal → 0, got {s}");
}
#[test]
fn combined_normalises_by_weight_sum() {
let weights = MaturityWeights {
recency: 0.0,
hit_success: 1.0,
edge_degree: 0.0,
redundancy: 0.0,
};
let b = MaturityBreakdown {
recency: 0.0,
hit_success: 0.9,
edge_degree: 0.0,
redundancy: 0.0,
combined: 0.0,
};
let c = combined_score(weights, b);
assert!((c - 0.9).abs() < 1e-5, "expected 0.9, got {c}");
}
#[test]
fn combined_all_zero_weights_is_zero() {
let weights = MaturityWeights {
recency: 0.0,
hit_success: 0.0,
edge_degree: 0.0,
redundancy: 0.0,
};
let b = MaturityBreakdown {
recency: 1.0,
hit_success: 1.0,
edge_degree: 1.0,
redundancy: 1.0,
combined: 0.0,
};
assert_eq!(combined_score(weights, b), 0.0);
}
#[tokio::test]
async fn edge_degree_component_zero_when_no_relations() {
let storage = Arc::new(DuckDbStorage::open_in_memory().unwrap());
let index = Arc::new(UsearchIndex::new(3).unwrap());
let embedding = Arc::new(NoopEmbedding::new(3));
let engine = MnemoEngine::new(storage, index, embedding, "a".to_string(), None);
let r = record(chrono::Utc::now(), 0, 0);
engine.storage.insert_memory(&r).await.unwrap();
let score = edge_degree_component(&engine, &[&r], 6.0).await.unwrap();
assert_eq!(score, 0.0);
}
}