Include identity nodes in memory scoring

Identity memory nodes now participate in importance scoring alongside
conversation memories. Score loading/saving handles both sections, and
the conscious screen uses node.label() consistently for memory display.

Co-Authored-By: Proof of Concept <poc@bcachefs.org>

src/mind/mod.rs

@@ -33,6 +33,36 @@ pub use unconscious::{UnconsciousSnapshot, Unconscious};
 
 use crate::agent::context::{AstNode, NodeBody, Section, Ast, ContextState};
 
+fn match_scores(
+    nodes: &[AstNode],
+    scores: &std::collections::BTreeMap<String, f64>,
+) -> Vec<(usize, f64)> {
+    nodes.iter().enumerate()
+        .filter_map(|(i, node)| {
+            if let AstNode::Leaf(leaf) = node {
+                if let NodeBody::Memory { key, .. } = leaf.body() {
+                    return scores.get(key.as_str()).map(|&s| (i, s));
+                }
+            }
+            None
+        }).collect()
+}
+
+fn find_memory_by_key(ctx: &ContextState, key: &str) -> Option<(Section, usize)> {
+    [(Section::Identity, ctx.identity()), (Section::Conversation, ctx.conversation())]
+        .into_iter()
+        .find_map(|(section, nodes)| {
+            nodes.iter().enumerate().find_map(|(i, node)| {
+                if let AstNode::Leaf(leaf) = node {
+                    if let NodeBody::Memory { key: k, .. } = leaf.body() {
+                        if k == key { return Some((section, i)); }
+                    }
+                }
+                None
+            })
+        })
+}
+
 fn load_memory_scores(ctx: &mut ContextState, path: &std::path::Path) {
     let data = match std::fs::read_to_string(path) {
         Ok(d) => d,
@@ -42,25 +72,24 @@ fn load_memory_scores(ctx: &mut ContextState, path: &std::path::Path) {
         Ok(s) => s,
         Err(_) => return,
     };
-    let mut applied = 0;
+    let identity_scores = match_scores(ctx.identity(), &scores);
-    for i in 0..ctx.conversation().len() {
+    let conv_scores = match_scores(ctx.conversation(), &scores);
-        if let AstNode::Leaf(leaf) = &ctx.conversation()[i] {
+    let applied = identity_scores.len() + conv_scores.len();
-            if let NodeBody::Memory { key, .. } = leaf.body() {
+    for (i, s) in identity_scores {
-                if let Some(&s) = scores.get(key.as_str()) {
+        ctx.set_score(Section::Identity, i, Some(s));
+    }
+    for (i, s) in conv_scores {
         ctx.set_score(Section::Conversation, i, Some(s));
-                    applied += 1;
-                }
-            }
-        }
     }
     if applied > 0 {
         dbglog!("[scoring] loaded {} scores from {}", applied, path.display());
     }
 }
 
-/// Collect scored memory keys from conversation entries.
+/// Collect scored memory keys from identity and conversation entries.
 fn collect_memory_scores(ctx: &ContextState) -> std::collections::BTreeMap<String, f64> {
-    ctx.conversation().iter()
+    ctx.identity().iter()
+        .chain(ctx.conversation().iter())
         .filter_map(|node| {
             if let AstNode::Leaf(leaf) = node {
                 if let NodeBody::Memory { key, score: Some(s), .. } = leaf.body() {
@@ -531,14 +560,10 @@ impl Mind {
                     async move {
                         let scores_snapshot = {
                             let mut ctx = agent.context.lock().await;
-                            for i in 0..ctx.conversation().len() {
+                            // Find memory by key in identity or conversation
-                                if let AstNode::Leaf(leaf) = &ctx.conversation()[i] {
+                            let found = find_memory_by_key(&ctx, &key);
-                                    if let NodeBody::Memory { key: k, .. } = leaf.body() {
+                            if let Some((section, i)) = found {
-                                        if *k == key {
+                                ctx.set_score(section, i, Some(score));
-                                            ctx.set_score(Section::Conversation, i, Some(score));
-                                        }
-                                    }
-                                }
                             }
                             let snapshot = collect_memory_scores(&ctx);
                             drop(ctx);

src/subconscious/learn.rs

@@ -62,9 +62,17 @@ fn build_token_ids(
     for node in context.system() {
         ids.extend(node.token_ids());
     }
+    // Identity nodes can be filtered by key for scoring
     for node in context.identity() {
+        let skip = match &filter {
+            Filter::SkipKey(key) => memory_key(node) == Some(*key),
+            Filter::SkipAllMemories => is_memory(node),
+            _ => false,
+        };
+        if !skip {
             ids.extend(node.token_ids());
         }
+    }
     for node in context.journal() {
         ids.extend(node.token_ids());
     }
@@ -175,7 +183,9 @@ pub async fn score_memories(
     // Collect memory keys and response indices under a brief lock
     let (memory_keys, response_indices) = {
         let ctx = agent.context.lock().await;
-        let mut keys: Vec<String> = ctx.conversation().iter()
+        // Include identity nodes and conversation memories
+        let mut keys: Vec<String> = ctx.identity().iter()
+            .chain(ctx.conversation().iter())
             .filter_map(|node| memory_key(node).map(String::from))
             .collect();
         keys.dedup();
@@ -331,7 +341,10 @@ where
 
     {
         let store = &*store_arc;
-        for (i, node) in context.conversation().iter().enumerate() {
+        // Identity nodes always score at position 0; conversation nodes at their index
+        let identity_nodes = context.identity().iter().map(|n| (0, n));
+        let conv_nodes = context.conversation().iter().enumerate();
+        for (pos, node) in identity_nodes.chain(conv_nodes) {
             if let Some(key) = memory_key(node) {
                 if !seen.insert(key.to_owned()) { continue; }
                 let last_scored = store.get_node(key)
@@ -340,7 +353,7 @@ where
                     .map(|n| n.last_scored)
                     .unwrap_or(0);
                 if now - last_scored >= max_age_secs {
-                    candidates.push((i, key.to_owned(), last_scored));
+                    candidates.push((pos, key.to_owned(), last_scored));
                 }
             }
         }

src/user/context.rs

@@ -37,17 +37,14 @@ impl ConsciousScreen {
         let mut unscored = 0usize;
         for node in ctx.conversation() {
             if let AstNode::Leaf(leaf) = node {
-                if let NodeBody::Memory { key, score, text } = leaf.body() {
+                if let NodeBody::Memory { score, text, .. } = leaf.body() {
-                    let status = match score {
+                    if score.is_some() { scored += 1; } else { unscored += 1; }
-                        Some(s) => { scored += 1; format!("{:.2}", s) }
-                        None => { unscored += 1; String::new() }
-                    };
                     mem_children.push(SectionView {
-                        name: key.clone(),
+                        name: node.label(),
                         tokens: node.tokens(),
                         content: text.clone(),
                         children: Vec::new(),
-                        status,
+                        status: String::new(),
                     });
                 }
             }