From 19205b9bae1af307767f38bdb751174b8ac154d2 Mon Sep 17 00:00:00 2001
From: Kent Overstreet <kent.overstreet@linux.dev>
Date: Thu, 2 Apr 2026 22:27:43 -0400
Subject: [PATCH] show scoring progress and per-response memory attribution
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Status bar shows "scoring 3/7..." during scoring. Debug pane logs
per-memory importance and top-5 response breakdowns. F10 context
screen shows which memories were important for each assistant
response as drilldown children (← memory_key (score)).

Added important_memories_for_entry() to look up the matrix by
conversation entry index.

Co-Authored-By: Proof of Concept <poc@bcachefs.org>
---
 src/agent/runner.rs     | 18 ++++++++++++-
 src/thought/training.rs | 60 ++++++++++++++++++++++++++++++++++++-----
 2 files changed, 71 insertions(+), 7 deletions(-)

diff --git a/src/agent/runner.rs b/src/agent/runner.rs
index 992e6bb..a17756b 100644
--- a/src/agent/runner.rs
+++ b/src/agent/runner.rs
@@ -739,11 +739,27 @@ impl Agent {
                         Role::System => "system".to_string(),
                     }
                 };
+                // Show which memories were important for this response
+                let children = if m.role == Role::Assistant {
+                    self.memory_scores.as_ref()
+                        .map(|s| s.important_memories_for_entry(i))
+                        .unwrap_or_default()
+                        .into_iter()
+                        .map(|(key, score)| ContextSection {
+                            name: format!("← {} ({:.1})", key, score),
+                            tokens: 0,
+                            content: String::new(),
+                            children: Vec::new(),
+                        })
+                        .collect()
+                } else {
+                    Vec::new()
+                };
                 ContextSection {
                     name: format!("[{}] {}: {}", i, role_name, label),
                     tokens,
                     content: text,
-                    children: Vec::new(),
+                    children,
                 }
             })
             .collect();
diff --git a/src/thought/training.rs b/src/thought/training.rs
index e1fdfec..2cfe73a 100644
--- a/src/thought/training.rs
+++ b/src/thought/training.rs
@@ -21,6 +21,27 @@ pub struct MemoryScore {
     pub matrix: Vec<Vec<f64>>,
     /// Keys of memories that were scored
     pub memory_keys: Vec<String>,
+    /// Conversation entry indices of the assistant responses (maps response_idx → entry_idx)
+    pub response_entry_indices: Vec<usize>,
+}
+
+impl MemoryScore {
+    /// Get the most important memories for a given conversation entry index.
+    /// Returns (memory_key, divergence_score) sorted by importance.
+    pub fn important_memories_for_entry(&self, entry_idx: usize) -> Vec<(&str, f64)> {
+        let Some(resp_idx) = self.response_entry_indices.iter().position(|&i| i == entry_idx)
+        else { return Vec::new() };
+
+        let mut result: Vec<(&str, f64)> = self.memory_keys.iter()
+            .zip(self.matrix.iter())
+            .filter_map(|(key, row)| {
+                let score = row.get(resp_idx).copied().unwrap_or(0.0);
+                if score > 0.01 { Some((key.as_str(), score)) } else { None }
+            })
+            .collect();
+        result.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
+        result
+    }
 }
 
 /// Score how important each memory is to the conversation.
@@ -53,6 +74,7 @@ pub async fn score_memories(
             response_scores: Vec::new(),
             matrix: Vec::new(),
             memory_keys: Vec::new(),
+            response_entry_indices: Vec::new(),
         });
     }
 
@@ -74,6 +96,9 @@ pub async fn score_memories(
     let memory_keys: Vec<String> = memories.iter().map(|(_, k)| k.clone()).collect();
 
     for (mem_idx, (entry_idx, key)) in memories.iter().enumerate() {
+        let _ = ui_tx.send(UiMessage::Activity(format!(
+            "scoring {}/{}...", mem_idx + 1, memories.len(),
+        )));
         let _ = ui_tx.send(UiMessage::Debug(format!(
             "[training] scoring memory {}/{}: {}",
             mem_idx + 1, memories.len(), key,
@@ -117,18 +142,41 @@ pub async fn score_memories(
         }
     }
 
-    let _ = ui_tx.send(UiMessage::Debug(format!(
-        "[training] done. top memory: {:?}",
-        memory_weights.iter()
-            .max_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal))
-            .map(|(k, v)| format!("{}: {:.1}", k, v)),
-    )));
+    let _ = ui_tx.send(UiMessage::Activity(String::new()));
+
+    // Log summary per memory
+    for (key, score) in &memory_weights {
+        let _ = ui_tx.send(UiMessage::Debug(format!(
+            "[training] {} → importance {:.1}", key, score,
+        )));
+    }
+
+    // Log per-response breakdown for the most important memories
+    let mut sorted_mems: Vec<(usize, &str, f64)> = memory_keys.iter().enumerate()
+        .map(|(i, k)| (i, k.as_str(), memory_weights[i].1))
+        .collect();
+    sorted_mems.sort_by(|a, b| b.2.partial_cmp(&a.2).unwrap_or(std::cmp::Ordering::Equal));
+
+    for (mem_i, key, total) in sorted_mems.iter().take(5) {
+        if *total <= 0.0 { continue; }
+        let row = &matrix[*mem_i];
+        let top_responses: Vec<String> = row.iter().enumerate()
+            .filter(|(_, v)| **v > 0.1)
+            .map(|(j, v)| format!("resp[{}]={:.1}", j, v))
+            .collect();
+        if !top_responses.is_empty() {
+            let _ = ui_tx.send(UiMessage::Debug(format!(
+                "[training]   {} ({:.1}): {}", key, total, top_responses.join(", "),
+            )));
+        }
+    }
 
     Ok(MemoryScore {
         memory_weights,
         response_scores,
         matrix,
         memory_keys,
+        response_entry_indices: response_indices,
     })
 }