use std::collections::{HashMap, HashSet};
use mentedb_core::error::MenteResult;
use mentedb_core::memory::{AttributeValue, MemoryNode, MemoryType};
use mentedb_core::types::{AgentId, MemoryId, UserId};
use crate::MenteDb;
pub const ATTR_INJECTION_SHOWN: &str = "injection_shown";
pub const ATTR_INJECTION_USED: &str = "injection_used";
#[derive(Debug, Clone)]
pub struct InjectionConfig {
pub candidate_pool: usize,
pub cluster_dominant_fraction: f32,
pub cluster_fill_max: usize,
pub cluster_max_span: usize,
pub use_lexical_leg: bool,
pub own_agent_boost: f32,
pub mmr_lambda: f32,
pub knee_gap_ratio: f32,
pub demotion_shown_min: i64,
pub demotion_factor: f32,
pub demotion_drop_shown: i64,
pub tail_min_keep: usize,
pub tail_floor_fraction: f32,
pub cross_project_episodic: bool,
pub cross_project_floor_fraction: f32,
pub used_similarity: f32,
pub use_reinforcement: f32,
pub graph_expansion_max: usize,
pub graph_expansion_decay: f32,
pub mode_trigger_gate: f32,
pub mode_rules_max: usize,
}
impl Default for InjectionConfig {
fn default() -> Self {
Self {
candidate_pool: 48,
cluster_dominant_fraction: 0.0,
cluster_fill_max: 4,
cluster_max_span: 12,
use_lexical_leg: true,
own_agent_boost: 1.25,
mmr_lambda: 0.7,
knee_gap_ratio: 2.0,
demotion_shown_min: 5,
demotion_factor: 0.5,
demotion_drop_shown: 12,
tail_min_keep: 3,
tail_floor_fraction: 0.55,
cross_project_episodic: false,
cross_project_floor_fraction: 0.8,
used_similarity: 0.6,
use_reinforcement: 0.05,
graph_expansion_max: 0,
graph_expansion_decay: 0.5,
mode_trigger_gate: 0.40,
mode_rules_max: 4,
}
}
}
pub struct InjectionQuery<'a> {
pub embedding: &'a [f32],
pub query_text: Option<&'a str>,
pub session_id: Option<&'a str>,
pub exclude_ids: &'a [MemoryId],
pub max_items: usize,
pub max_episodic: usize,
pub agent_id: Option<AgentId>,
pub user_id: Option<UserId>,
pub current_project: Option<&'a str>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SelectionReason {
Pinned,
Relevant,
ModeRule,
}
pub struct InjectionCandidate {
pub node: MemoryNode,
pub score: f32,
pub reason: SelectionReason,
}
fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
if a.len() != b.len() || a.is_empty() {
return 0.0;
}
let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
let na: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
let nb: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
if na == 0.0 || nb == 0.0 {
return 0.0;
}
dot / (na * nb)
}
fn attr_count(node: &MemoryNode, key: &str) -> i64 {
match node.attributes.get(key) {
Some(AttributeValue::Integer(n)) => *n,
_ => 0,
}
}
fn bump_attr(node: &mut MemoryNode, key: &str) {
let next = attr_count(node, key) + 1;
node.attributes
.insert(key.to_string(), AttributeValue::Integer(next));
}
fn has_tag(node: &MemoryNode, tag: &str) -> bool {
node.tags.iter().any(|t| t == tag)
}
fn now_us() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_micros() as u64
}
fn knee_cutoff(scores: &[f32], gap_ratio: f32) -> usize {
if scores.len() <= 1 {
return scores.len();
}
let mut best_ratio = 0.0f32;
let mut cut = scores.len();
for i in 0..scores.len() - 1 {
let hi = scores[i];
let lo = scores[i + 1].max(f32::EPSILON);
let ratio = hi / lo;
if ratio > best_ratio {
best_ratio = ratio;
cut = i + 1;
}
}
if best_ratio >= gap_ratio {
cut
} else {
scores.len()
}
}
fn project_of(node: &MemoryNode) -> Option<&str> {
node.tags
.iter()
.find_map(|t| t.strip_prefix("scope:project:"))
}
fn cross_project_allowed(
node: &MemoryNode,
cos: f32,
best_cos: f32,
current_project: Option<&str>,
cfg: &InjectionConfig,
) -> bool {
let Some(current) = current_project else {
return true;
};
let Some(mem_project) = project_of(node) else {
return true;
};
if mem_project == current {
return true;
}
if node.memory_type == MemoryType::Episodic {
return cfg.cross_project_episodic;
}
cos >= best_cos * cfg.cross_project_floor_fraction
}
fn passes_tail_gate(idx: usize, cos: f32, best_cos: f32, cfg: &InjectionConfig) -> bool {
if idx < cfg.tail_min_keep {
return true;
}
cos >= best_cos * cfg.tail_floor_fraction
}
impl MenteDb {
pub fn recall_for_injection(
&self,
query: &InjectionQuery<'_>,
) -> MenteResult<Vec<InjectionCandidate>> {
let cfg = self.cognitive_config.injection_config.clone();
let excluded: HashSet<MemoryId> = query.exclude_ids.iter().copied().collect();
let session_tag = query.session_id.map(|s| format!("session:{s}"));
let hits = self
.recall_hybrid_scoped_at_mode(
query.embedding,
query.query_text.filter(|_| cfg.use_lexical_leg),
cfg.candidate_pool,
now_us(),
None,
false,
None,
query.agent_id,
query.user_id,
query.current_project,
)
.unwrap_or_default();
let mut candidates: Vec<(MemoryNode, f32, f32)> = Vec::new();
for (id, score) in hits {
if excluded.contains(&id) {
continue;
}
let Ok(node) = self.get_memory(id) else {
continue;
};
if has_tag(&node, "action")
|| has_tag(&node, "scope:always")
|| has_tag(&node, "ghost-memory")
|| node.tags.iter().any(|t| t.starts_with("mode-exemplar:"))
{
continue;
}
if let Some(ref st) = session_tag
&& has_tag(&node, st)
{
continue;
}
let shown = attr_count(&node, ATTR_INJECTION_SHOWN);
let used = attr_count(&node, ATTR_INJECTION_USED);
if shown >= cfg.demotion_drop_shown && used == 0 {
continue;
}
let mut adjusted = if shown >= cfg.demotion_shown_min && used == 0 {
score * cfg.demotion_factor
} else {
score
};
if let Some(agent) = query.agent_id
&& node.agent_id == agent
&& !agent.0.is_nil()
{
adjusted *= cfg.own_agent_boost;
}
let cos = cosine_similarity(query.embedding, &node.embedding);
candidates.push((node, adjusted, cos));
}
let cluster_pool: Vec<(MemoryNode, f32)> = if cfg.cluster_dominant_fraction > 0.0 {
candidates
.iter()
.filter(|(node, _, _)| node.tags.iter().any(|t| t.starts_with("section:")))
.map(|(node, _, cos)| (node.clone(), *cos))
.collect()
} else {
Vec::new()
};
let best_cos = candidates.iter().map(|(_, _, c)| *c).fold(0.0f32, f32::max);
candidates.retain(|(node, _, cos)| {
cross_project_allowed(node, *cos, best_cos, query.current_project, &cfg)
});
candidates.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
let scores: Vec<f32> = candidates.iter().map(|(_, s, _)| *s).collect();
candidates.truncate(knee_cutoff(&scores, cfg.knee_gap_ratio));
let mut idx = 0usize;
candidates.retain(|(_, _, cos)| {
let keep = passes_tail_gate(idx, *cos, best_cos, &cfg);
idx += 1;
keep
});
let mut scored: Vec<(MemoryNode, f32)> =
candidates.into_iter().map(|(n, s, _)| (n, s)).collect();
if cfg.graph_expansion_max > 0 && !scored.is_empty() {
let now = now_us();
let mut present: HashSet<MemoryId> = scored.iter().map(|(n, _)| n.id).collect();
present.extend(excluded.iter().copied());
let seeds: Vec<(MemoryId, f32)> = scored.iter().map(|(n, s)| (n.id, *s)).collect();
let g = self.graph.graph();
let mut added = 0usize;
'seeds: for (seed_id, seed_score) in seeds {
let mut edges = g.outgoing(seed_id);
edges.extend(g.incoming(seed_id));
for (nbr_id, edge) in edges {
if added >= cfg.graph_expansion_max {
break 'seeds;
}
if present.contains(&nbr_id) {
continue;
}
let Ok(node) = self.get_memory(nbr_id) else {
continue;
};
if has_tag(&node, "action")
|| has_tag(&node, "scope:always")
|| has_tag(&node, "ghost-memory")
|| !node.is_valid_at(now)
|| !crate::agent_visible(node.agent_id, query.agent_id)
|| !crate::user_visible(node.user_id, query.user_id)
{
continue;
}
if let Some(ref st) = session_tag
&& has_tag(&node, st)
{
continue;
}
let score = seed_score * edge.weight * cfg.graph_expansion_decay;
present.insert(nbr_id);
scored.push((node, score));
added += 1;
}
}
}
let dominant_section: Option<String> = if cfg.cluster_dominant_fraction > 0.0 {
let head = scored.len().min(query.max_items.max(3));
let mut counts: HashMap<String, usize> = HashMap::new();
for (node, _) in &scored[..head] {
for tag in &node.tags {
if let Some(sec) = tag.strip_prefix("section:") {
*counts.entry(sec.to_string()).or_insert(0) += 1;
}
}
}
let mut span: HashMap<&str, usize> = HashMap::new();
for (node, _) in &cluster_pool {
for tag in &node.tags {
if let Some(sec) = tag.strip_prefix("section:") {
*span.entry(sec).or_insert(0) += 1;
}
}
}
let need = ((cfg.cluster_dominant_fraction * head as f32).ceil() as usize).max(2);
counts
.into_iter()
.filter(|(sec, n)| {
*n >= need
&& span.get(sec.as_str()).copied().unwrap_or(0) <= cfg.cluster_max_span
})
.max_by_key(|(_, n)| *n)
.map(|(sec, _)| sec)
} else {
None
};
let top_score = scored
.first()
.map(|(_, s)| *s)
.unwrap_or(1.0)
.max(f32::EPSILON);
let mut remaining: Vec<(MemoryNode, f32)> = scored;
let mut selected: Vec<InjectionCandidate> = Vec::new();
let mut episodic_count = 0usize;
while selected.len() < query.max_items && !remaining.is_empty() {
let mut best_idx: Option<usize> = None;
let mut best_value = f32::NEG_INFINITY;
for (idx, (node, score)) in remaining.iter().enumerate() {
if node.memory_type == MemoryType::Episodic && episodic_count >= query.max_episodic
{
continue;
}
let redundancy = selected
.iter()
.map(|s| cosine_similarity(&node.embedding, &s.node.embedding))
.fold(0.0f32, f32::max);
let value =
cfg.mmr_lambda * (score / top_score) - (1.0 - cfg.mmr_lambda) * redundancy;
if value > best_value {
best_value = value;
best_idx = Some(idx);
}
}
let Some(idx) = best_idx else {
break;
};
let (node, score) = remaining.remove(idx);
if node.memory_type == MemoryType::Episodic {
episodic_count += 1;
}
selected.push(InjectionCandidate {
node,
score,
reason: SelectionReason::Relevant,
});
}
if selected.len() >= 3
&& let Some(sec) = dominant_section
{
{
let sec_tag = format!("section:{sec}");
let mut siblings: Vec<(MemoryNode, f32)> = cluster_pool
.into_iter()
.filter(|(node, _)| {
node.tags.iter().any(|t| t == &sec_tag)
&& !selected.iter().any(|c| c.node.id == node.id)
})
.collect();
siblings.sort_by(|a, b| b.1.total_cmp(&a.1));
siblings.truncate(cfg.cluster_fill_max);
let mut evict: Vec<usize> = (1..selected.len())
.filter(|i| !selected[*i].node.tags.iter().any(|t| t == &sec_tag))
.collect();
evict.sort_by(|a, b| {
selected[*a]
.score
.partial_cmp(&selected[*b].score)
.unwrap_or(std::cmp::Ordering::Equal)
});
for (pos, (node, cos)) in evict.into_iter().zip(siblings) {
selected[pos] = InjectionCandidate {
node,
score: cos,
reason: SelectionReason::Relevant,
};
}
}
}
let mut result: Vec<InjectionCandidate> = Vec::new();
for mid in self.index.bitmap.query_tag("scope:always") {
if excluded.contains(&mid) {
continue;
}
if let Ok(node) = self.get_memory(mid)
&& crate::agent_visible(node.agent_id, query.agent_id)
&& crate::user_visible(node.user_id, query.user_id)
{
result.push(InjectionCandidate {
node,
score: 0.0,
reason: SelectionReason::Pinned,
});
}
}
if cfg.mode_trigger_gate > 0.0 && !query.embedding.is_empty() {
let now = now_us();
let mut mode_rules: Vec<(f32, MemoryNode)> = Vec::new();
for tag in self.index.bitmap.tags_with_prefix("trigger:") {
let trigger = tag.trim_start_matches("trigger:");
let exemplar_tag = format!("mode-exemplar:{trigger}");
let mut activation = f32::MIN;
let exemplar_ids = self.index.bitmap.query_tag(&exemplar_tag);
if exemplar_ids.is_empty() {
for mid in self.index.bitmap.query_tag(&tag) {
if let Ok(node) = self.get_memory(mid)
&& node.is_valid_at(now)
&& !node.embedding.is_empty()
&& crate::agent_visible(node.agent_id, query.agent_id)
&& crate::user_visible(node.user_id, query.user_id)
{
let sim = cosine_similarity(query.embedding, &node.embedding);
activation = activation.max(sim);
}
}
} else {
for mid in exemplar_ids {
if let Ok(node) = self.get_memory(mid)
&& node.is_valid_at(now)
&& !node.embedding.is_empty()
&& crate::agent_visible(node.agent_id, query.agent_id)
&& crate::user_visible(node.user_id, query.user_id)
{
let sim = cosine_similarity(query.embedding, &node.embedding);
activation = activation.max(sim);
}
}
}
if activation < cfg.mode_trigger_gate {
continue;
}
for mid in self.index.bitmap.query_tag(&tag) {
if excluded.contains(&mid)
|| result.iter().any(|c| c.node.id == mid)
|| selected.iter().any(|c| c.node.id == mid)
|| mode_rules.iter().any(|(_, n)| n.id == mid)
{
continue;
}
let Ok(node) = self.get_memory(mid) else {
continue;
};
if !node.tags.iter().any(|t| t == &tag)
|| !node.is_valid_at(now)
|| !crate::agent_visible(node.agent_id, query.agent_id)
|| !crate::user_visible(node.user_id, query.user_id)
{
continue;
}
mode_rules.push((activation, node));
}
}
mode_rules.sort_by(|a, b| b.0.total_cmp(&a.0));
mode_rules.truncate(cfg.mode_rules_max);
for (sim, node) in mode_rules {
result.push(InjectionCandidate {
node,
score: sim,
reason: SelectionReason::ModeRule,
});
}
}
result.extend(selected);
Ok(result)
}
pub fn record_injection_outcome(
&self,
shown: &[MemoryId],
reply_embedding: Option<&[f32]>,
) -> MenteResult<(usize, usize)> {
let cfg = self.cognitive_config.injection_config.clone();
let mut updated = 0usize;
let mut used_total = 0usize;
let now = now_us();
for id in shown {
let pid = {
let pm = self.page_map.read();
match pm.get(id) {
Some(p) => *p,
None => continue,
}
};
let Ok(mut node) = self.storage.load_memory(pid) else {
continue;
};
bump_attr(&mut node, ATTR_INJECTION_SHOWN);
node.access_count = node.access_count.saturating_add(1);
node.accessed_at = now;
let used = reply_embedding
.map(|re| cosine_similarity(&node.embedding, re) >= cfg.used_similarity)
.unwrap_or(false);
if used {
bump_attr(&mut node, ATTR_INJECTION_USED);
node.salience = (node.salience + cfg.use_reinforcement).min(1.0);
used_total += 1;
}
self.storage.update_memory(pid, &node)?;
updated += 1;
}
Ok((updated, used_total))
}
}
#[cfg(test)]
mod tests {
use super::*;
use mentedb_core::types::AgentId;
use uuid::Uuid;
#[test]
fn knee_cuts_at_largest_gap() {
let ratio = InjectionConfig::default().knee_gap_ratio;
assert_eq!(knee_cutoff(&[1.0, 0.9, 0.8, 0.1, 0.09], ratio), 3);
assert_eq!(knee_cutoff(&[1.0, 0.9, 0.8, 0.7], ratio), 4);
assert_eq!(knee_cutoff(&[1.0], ratio), 1);
assert_eq!(knee_cutoff(&[], ratio), 0);
}
#[test]
fn attr_counters_roundtrip() {
let mut node = MemoryNode::new(
AgentId(Uuid::nil()),
MemoryType::Semantic,
"x".into(),
vec![0.0; 4],
);
assert_eq!(attr_count(&node, ATTR_INJECTION_SHOWN), 0);
bump_attr(&mut node, ATTR_INJECTION_SHOWN);
bump_attr(&mut node, ATTR_INJECTION_SHOWN);
assert_eq!(attr_count(&node, ATTR_INJECTION_SHOWN), 2);
}
fn node_with_tags(memory_type: MemoryType, tags: &[&str]) -> MemoryNode {
let mut node = MemoryNode::new(AgentId(Uuid::nil()), memory_type, "x".into(), vec![0.0; 4]);
node.tags = tags.iter().map(|t| t.to_string()).collect();
node
}
#[test]
fn tail_gate_trusts_head_and_prunes_weak_tail() {
let cfg = InjectionConfig::default();
let best = 0.8;
assert!(passes_tail_gate(0, 0.05, best, &cfg));
assert!(passes_tail_gate(2, 0.05, best, &cfg));
assert!(passes_tail_gate(3, 0.6, best, &cfg));
assert!(!passes_tail_gate(3, 0.3, best, &cfg));
assert!(passes_tail_gate(5, 0.0, 0.0, &cfg));
}
#[test]
fn cross_project_policy() {
let cfg = InjectionConfig::default();
let best = 0.8;
let foreign_epi = node_with_tags(MemoryType::Episodic, &["scope:project:other"]);
let foreign_sem = node_with_tags(MemoryType::Semantic, &["scope:project:other"]);
let same_epi = node_with_tags(MemoryType::Episodic, &["scope:project:mine"]);
let global_sem = node_with_tags(MemoryType::Semantic, &[]);
assert!(cross_project_allowed(&foreign_epi, 0.1, best, None, &cfg));
assert!(cross_project_allowed(
&same_epi,
0.1,
best,
Some("mine"),
&cfg
));
assert!(cross_project_allowed(
&global_sem,
0.1,
best,
Some("mine"),
&cfg
));
assert!(!cross_project_allowed(
&foreign_epi,
0.79,
best,
Some("mine"),
&cfg
));
assert!(cross_project_allowed(
&foreign_sem,
0.7,
best,
Some("mine"),
&cfg
));
assert!(!cross_project_allowed(
&foreign_sem,
0.5,
best,
Some("mine"),
&cfg
));
}
#[test]
fn project_of_reads_scope_tag() {
let node = node_with_tags(MemoryType::Semantic, &["session:s1", "scope:project:apex"]);
assert_eq!(project_of(&node), Some("apex"));
let untagged = node_with_tags(MemoryType::Semantic, &["session:s1"]);
assert_eq!(project_of(&untagged), None);
}
}