use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::time::Instant;
use super::types::MemoryId;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RelevanceWeights {
pub semantic: f32,
pub recency: f32,
pub strength: f32,
#[serde(default)]
pub entity_overlap: f32,
#[serde(default)]
pub type_boost: f32,
#[serde(default)]
pub file_match: f32,
#[serde(default)]
pub suppression: f32,
#[serde(default)]
pub episode_coherence: f32,
#[serde(default)]
pub graph_activation: f32,
#[serde(default)]
pub linguistic_score: f32,
}
impl Default for RelevanceWeights {
fn default() -> Self {
Self {
semantic: 0.40, recency: 0.08, strength: 0.08, entity_overlap: 0.08, type_boost: 0.06, file_match: 0.04, suppression: 0.02, episode_coherence: 0.06, graph_activation: 0.10, linguistic_score: 0.08, }
}
}
impl RelevanceWeights {
pub fn legacy() -> Self {
Self {
semantic: 0.5,
recency: 0.3,
strength: 0.2,
entity_overlap: 0.0,
type_boost: 0.0,
file_match: 0.0,
suppression: 0.0,
episode_coherence: 0.0,
graph_activation: 0.0,
linguistic_score: 0.0,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InjectionConfig {
pub min_relevance: f32,
pub max_per_message: usize,
pub cooldown_seconds: u64,
pub weights: RelevanceWeights,
pub recency_decay_rate: f32,
}
impl Default for InjectionConfig {
fn default() -> Self {
Self {
min_relevance: 0.50, max_per_message: 3,
cooldown_seconds: 180,
weights: RelevanceWeights::default(),
recency_decay_rate: 0.01,
}
}
}
impl InjectionConfig {
pub fn legacy() -> Self {
Self {
min_relevance: 0.70,
max_per_message: 3,
cooldown_seconds: 180,
weights: RelevanceWeights::legacy(),
recency_decay_rate: 0.01,
}
}
}
#[derive(Debug, Clone)]
pub struct InjectionCandidate {
pub memory_id: MemoryId,
pub relevance_score: f32,
}
pub struct InjectionEngine {
config: InjectionConfig,
cooldowns: HashMap<MemoryId, Instant>,
}
impl InjectionEngine {
pub fn new(config: InjectionConfig) -> Self {
Self {
config,
cooldowns: HashMap::new(),
}
}
pub fn with_default_config() -> Self {
Self::new(InjectionConfig::default())
}
fn on_cooldown(&self, memory_id: &MemoryId) -> bool {
if let Some(last) = self.cooldowns.get(memory_id) {
last.elapsed().as_secs() < self.config.cooldown_seconds
} else {
false
}
}
pub fn select_for_injection(
&mut self,
mut candidates: Vec<InjectionCandidate>,
) -> Vec<MemoryId> {
candidates.sort_by(|a, b| b.relevance_score.total_cmp(&a.relevance_score));
let selected: Vec<MemoryId> = candidates
.into_iter()
.filter(|c| {
c.relevance_score >= self.config.min_relevance && !self.on_cooldown(&c.memory_id)
})
.take(self.config.max_per_message)
.map(|c| c.memory_id)
.collect();
let now = Instant::now();
for id in &selected {
self.cooldowns.insert(id.clone(), now);
}
selected
}
pub fn cleanup_cooldowns(&mut self) {
let threshold = self.config.cooldown_seconds;
self.cooldowns
.retain(|_, last| last.elapsed().as_secs() < threshold * 2);
}
pub fn config(&self) -> &InjectionConfig {
&self.config
}
pub fn set_config(&mut self, config: InjectionConfig) {
self.config = config;
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InjectionRecord {
pub memory_id: MemoryId,
pub injected_at: DateTime<Utc>,
pub relevance_score: f32,
pub context_signature: u64,
}
#[derive(Debug, Default)]
pub struct InjectionTracker {
pending: Vec<InjectionRecord>,
max_pending: usize,
}
impl InjectionTracker {
pub fn new(max_pending: usize) -> Self {
Self {
pending: Vec::new(),
max_pending,
}
}
pub fn record_injection(
&mut self,
memory_id: MemoryId,
relevance_score: f32,
context_signature: u64,
) {
let record = InjectionRecord {
memory_id,
injected_at: Utc::now(),
relevance_score,
context_signature,
};
self.pending.push(record);
if self.pending.len() > self.max_pending {
self.pending.remove(0);
}
}
pub fn pending_injections(&self) -> &[InjectionRecord] {
&self.pending
}
pub fn clear_old(&mut self, max_age_seconds: i64) {
let cutoff = Utc::now() - chrono::Duration::seconds(max_age_seconds);
self.pending.retain(|r| r.injected_at > cutoff);
}
pub fn mark_processed(&mut self, memory_id: &MemoryId) {
self.pending.retain(|r| &r.memory_id != memory_id);
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FeedbackSignal {
Positive,
Negative,
Neutral,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UserInjectionProfile {
pub user_id: String,
pub effective_threshold: f32,
pub positive_signals: u32,
pub negative_signals: u32,
pub updated_at: DateTime<Utc>,
}
impl UserInjectionProfile {
pub fn new(user_id: String) -> Self {
Self {
user_id,
effective_threshold: InjectionConfig::default().min_relevance,
positive_signals: 0,
negative_signals: 0,
updated_at: Utc::now(),
}
}
pub fn adjust(&mut self, signal: FeedbackSignal) {
match signal {
FeedbackSignal::Positive => {
self.positive_signals += 1;
self.effective_threshold = (self.effective_threshold - 0.01).max(0.50);
}
FeedbackSignal::Negative => {
self.negative_signals += 1;
self.effective_threshold = (self.effective_threshold + 0.02).min(0.90);
}
FeedbackSignal::Neutral => {}
}
self.updated_at = Utc::now();
}
pub fn signal_ratio(&self) -> f32 {
let total = self.positive_signals + self.negative_signals;
if total == 0 {
0.5 } else {
self.positive_signals as f32 / total as f32
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use uuid::Uuid;
#[test]
fn test_injection_engine_filtering() {
let mut engine = InjectionEngine::with_default_config();
let candidates = vec![
InjectionCandidate {
memory_id: MemoryId(Uuid::new_v4()),
relevance_score: 0.85,
},
InjectionCandidate {
memory_id: MemoryId(Uuid::new_v4()),
relevance_score: 0.45, },
InjectionCandidate {
memory_id: MemoryId(Uuid::new_v4()),
relevance_score: 0.75,
},
];
let selected = engine.select_for_injection(candidates);
assert_eq!(selected.len(), 2); }
#[test]
fn test_user_profile_adjustment() {
let mut profile = UserInjectionProfile::new("test-user".to_string());
assert_eq!(profile.effective_threshold, 0.50);
profile.adjust(FeedbackSignal::Positive);
assert!((profile.effective_threshold - 0.49).abs() < 0.01);
profile.adjust(FeedbackSignal::Negative);
assert!((profile.effective_threshold - 0.51).abs() < 0.01);
for _ in 0..20 {
profile.adjust(FeedbackSignal::Negative);
}
assert_eq!(profile.effective_threshold, 0.90);
}
}