kent/consciousness
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

cleanup: simplify MemoryNode, deduplicate tool dispatch

Browse source 
- Removed write/search/mark_used static methods from MemoryNode —
  those are store ops, not MemoryNode concerns
- Removed SearchResult duplicate — use query::engine::SearchResult
- Simplified Link to (String, f32) tuple — inline detection moved
  to render()
- Collapsed tool definitions to one-liners
- Consolidated store-mutation tools into with_store() helper
- Supersede uses store directly instead of MemoryNode round-trip

Co-Authored-By: Proof of Concept <poc@bcachefs.org>
This commit is contained in:
ProofOfConcept 2026-03-25 01:59:13 -04:00
parent 10932cb67e
commit 164a603c8e
2 changed files with 102 additions and 326 deletions
Download patch file Download diff file Expand all files Collapse all files

338
src/agent/tools/memory.rs
View file

@ -1,8 +1,6 @@
// tools/memory.rs — Native memory graph operations // tools/memory.rs — Native memory graph operations
// //
// Direct library calls into the store — no subprocess spawning. // Direct library calls into the store — no subprocess spawning.
// Returns MemoryNodes where possible so the agent can track what's
// loaded in its context window.
use anyhow::{Context, Result}; use anyhow::{Context, Result};
use serde_json::json; use serde_json::json;
@ -13,194 +11,66 @@ use crate::store::Store;
pub fn definitions() -> Vec<ToolDef> { pub fn definitions() -> Vec<ToolDef> {
vec![ vec![
ToolDef::new( ToolDef::new("memory_render",
"memory_render", "Read a memory node's content and links.",
"Read a memory node's content and links. Returns the full content \ json!({"type":"object","properties":{"key":{"type":"string","description":"Node key"}},"required":["key"]})),
with neighbor links sorted by strength.", ToolDef::new("memory_write",
json!({ "Create or update a memory node.",
"type": "object", json!({"type":"object","properties":{"key":{"type":"string","description":"Node key"},"content":{"type":"string","description":"Full content (markdown)"}},"required":["key","content"]})),
"properties": { ToolDef::new("memory_search",
"key": { "Search the memory graph by keyword.",
"type": "string", json!({"type":"object","properties":{"query":{"type":"string","description":"Search terms"}},"required":["query"]})),
"description": "Node key to render" ToolDef::new("memory_links",
}
},
"required": ["key"]
}),
),
ToolDef::new(
"memory_write",
"Create or update a memory node with new content. Use for writing \
prose, analysis, or any node content. Multi-line content is fine.",
json!({
"type": "object",
"properties": {
"key": {
"type": "string",
"description": "Node key to create or update"
},
"content": {
"type": "string",
"description": "Full content for the node (markdown)"
}
},
"required": ["key", "content"]
}),
),
ToolDef::new(
"memory_search",
"Search the memory graph for nodes by keyword.",
json!({
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "Search terms"
}
},
"required": ["query"]
}),
),
ToolDef::new(
"memory_links",
"Show a node's neighbors with link strengths.", "Show a node's neighbors with link strengths.",
json!({ json!({"type":"object","properties":{"key":{"type":"string","description":"Node key"}},"required":["key"]})),
"type": "object", ToolDef::new("memory_link_set",
"properties": { "Set link strength between two nodes.",
"key": { json!({"type":"object","properties":{"source":{"type":"string"},"target":{"type":"string"},"strength":{"type":"number","description":"0.01 to 1.0"}},"required":["source","target","strength"]})),
"type": "string", ToolDef::new("memory_link_add",
"description": "Node key to show links for"
}
},
"required": ["key"]
}),
),
ToolDef::new(
"memory_link_set",
"Set the strength of a link between two nodes. Also deduplicates \
if multiple links exist between the same pair.",
json!({
"type": "object",
"properties": {
"source": {
"type": "string",
"description": "Source node key"
},
"target": {
"type": "string",
"description": "Target node key"
},
"strength": {
"type": "number",
"description": "Link strength (0.01 to 1.0)"
}
},
"required": ["source", "target", "strength"]
}),
),
ToolDef::new(
"memory_link_add",
"Add a new link between two nodes.", "Add a new link between two nodes.",
json!({ json!({"type":"object","properties":{"source":{"type":"string"},"target":{"type":"string"}},"required":["source","target"]})),
"type": "object", ToolDef::new("memory_used",
"properties": { "Mark a node as useful (boosts weight).",
"source": { json!({"type":"object","properties":{"key":{"type":"string","description":"Node key"}},"required":["key"]})),
"type": "string", ToolDef::new("memory_weight_set",
"description": "Source node key" "Set a node's weight directly (0.01 to 1.0).",
}, json!({"type":"object","properties":{"key":{"type":"string"},"weight":{"type":"number","description":"0.01 to 1.0"}},"required":["key","weight"]})),
"target": { ToolDef::new("memory_supersede",
"type": "string", "Mark a node as superseded by another (sets weight to 0.01).",
"description": "Target node key" json!({"type":"object","properties":{"old_key":{"type":"string"},"new_key":{"type":"string"},"reason":{"type":"string"}},"required":["old_key","new_key"]})),
}
},
"required": ["source", "target"]
}),
),
ToolDef::new(
"memory_used",
"Mark a node as useful (boosts its weight in the graph).",
json!({
"type": "object",
"properties": {
"key": {
"type": "string",
"description": "Node key to mark as used"
}
},
"required": ["key"]
}),
),
ToolDef::new(
"memory_weight_set",
"Set a node's weight directly. Use to downweight junk nodes (0.01) \
or boost important ones. Normal range is 0.1 to 1.0.",
json!({
"type": "object",
"properties": {
"key": {
"type": "string",
"description": "Node key"
},
"weight": {
"type": "number",
"description": "New weight (0.01 to 1.0)"
}
},
"required": ["key", "weight"]
}),
),
ToolDef::new(
"memory_supersede",
"Mark a node as superseded by another. Sets the old node's weight \
to 0.01 and prepends a notice pointing to the replacement. Use \
when merging duplicates or replacing junk with proper content.",
json!({
"type": "object",
"properties": {
"old_key": {
"type": "string",
"description": "Node being superseded"
},
"new_key": {
"type": "string",
"description": "Replacement node"
},
"reason": {
"type": "string",
"description": "Why this node was superseded (e.g. 'merged into X', 'duplicate of Y')"
}
},
"required": ["old_key", "new_key"]
}),
),
] ]
} }
/// Dispatch a memory tool call. Direct library calls, no subprocesses. /// Dispatch a memory tool call. Direct library calls, no subprocesses.
pub fn dispatch(name: &str, args: &serde_json::Value, provenance: Option<&str>) -> Result<String> { pub fn dispatch(name: &str, args: &serde_json::Value, provenance: Option<&str>) -> Result<String> {
let prov = provenance.unwrap_or("manual"); let prov = provenance.unwrap_or("manual");
let result = match name { match name {
"memory_render" => { "memory_render" => {
let key = get_str(args, "key")?; let key = get_str(args, "key")?;
let node = MemoryNode::load(key) Ok(MemoryNode::load(key)
.ok_or_else(|| anyhow::anyhow!("node not found: {}", key))?; .ok_or_else(|| anyhow::anyhow!("node not found: {}", key))?
node.render() .render())
} }
"memory_write" => { "memory_write" => {
let key = get_str(args, "key")?; let key = get_str(args, "key")?;
let content = get_str(args, "content")?; let content = get_str(args, "content")?;
let node = MemoryNode::write(key, content, Some(prov)) let mut store = Store::load().map_err(|e| anyhow::anyhow!("{}", e))?;
let result = store.upsert_provenance(key, content, prov)
.map_err(|e| anyhow::anyhow!("{}", e))?; .map_err(|e| anyhow::anyhow!("{}", e))?;
format!("wrote '{}' (v{})", node.key, node.version) store.save().map_err(|e| anyhow::anyhow!("{}", e))?;
Ok(format!("{} '{}'", result, key))
} }
"memory_search" => { "memory_search" => {
let query = get_str(args, "query")?; let query = get_str(args, "query")?;
let results = MemoryNode::search(query) let store = Store::load().map_err(|e| anyhow::anyhow!("{}", e))?;
.map_err(|e| anyhow::anyhow!("{}", e))?; let results = crate::search::search(query, &store);
if results.is_empty() { if results.is_empty() {
"no results".to_string() Ok("no results".into())
} else { } else {
results.iter().map(|r| r.render()).collect::<Vec<_>>().join("\n") Ok(results.iter().take(20)
.map(|r| format!("({:.2}) {} — {}", r.activation, r.key,
r.snippet.as_deref().unwrap_or("")))
.collect::<Vec<_>>().join("\n"))
} }
} }
"memory_links" => { "memory_links" => {
@ -208,103 +78,75 @@ pub fn dispatch(name: &str, args: &serde_json::Value, provenance: Option<&str>)
let node = MemoryNode::load(key) let node = MemoryNode::load(key)
.ok_or_else(|| anyhow::anyhow!("node not found: {}", key))?; .ok_or_else(|| anyhow::anyhow!("node not found: {}", key))?;
let mut out = format!("Neighbors of '{}':\n", key); let mut out = format!("Neighbors of '{}':\n", key);
for link in &node.links { for (target, strength) in &node.links {
out.push_str(&format!(" ({:.2}) {}{}\n", out.push_str(&format!(" ({:.2}) {}\n", strength, target));
link.strength, link.target,
if link.inline { " [inline]" } else { "" }));
} }
out Ok(out)
} }
"memory_link_set" | "memory_link_add" | "memory_used" | "memory_weight_set" => {
with_store(name, args, prov)
}
"memory_supersede" => {
let old_key = get_str(args, "old_key")?;
let new_key = get_str(args, "new_key")?;
let reason = args.get("reason").and_then(|v| v.as_str()).unwrap_or("superseded");
let mut store = Store::load().map_err(|e| anyhow::anyhow!("{}", e))?;
let content = store.nodes.get(old_key)
.map(|n| n.content.clone())
.ok_or_else(|| anyhow::anyhow!("node not found: {}", old_key))?;
let notice = format!("**SUPERSEDED** by `{}` — {}\n\n---\n\n{}",
new_key, reason, content.trim());
store.upsert_provenance(old_key, &notice, prov)
.map_err(|e| anyhow::anyhow!("{}", e))?;
store.set_weight(old_key, 0.01).map_err(|e| anyhow::anyhow!("{}", e))?;
store.save().map_err(|e| anyhow::anyhow!("{}", e))?;
Ok(format!("superseded {} → {} ({})", old_key, new_key, reason))
}
_ => anyhow::bail!("Unknown memory tool: {}", name),
}
}
/// Store mutations that follow the same pattern: load, resolve, mutate, save.
fn with_store(name: &str, args: &serde_json::Value, prov: &str) -> Result<String> {
let mut store = Store::load().map_err(|e| anyhow::anyhow!("{}", e))?;
let msg = match name {
"memory_link_set" => { "memory_link_set" => {
let source = get_str(args, "source")?; let s = store.resolve_key(get_str(args, "source")?).map_err(|e| anyhow::anyhow!("{}", e))?;
let target = get_str(args, "target")?; let t = store.resolve_key(get_str(args, "target")?).map_err(|e| anyhow::anyhow!("{}", e))?;
let strength = get_f64(args, "strength")? as f32; let strength = get_f64(args, "strength")? as f32;
link_set(source, target, strength)? let old = store.set_link_strength(&s, &t, strength).map_err(|e| anyhow::anyhow!("{}", e))?;
format!("{} ↔ {} strength {:.2} → {:.2}", s, t, old, strength)
} }
"memory_link_add" => { "memory_link_add" => {
let source = get_str(args, "source")?; let s = store.resolve_key(get_str(args, "source")?).map_err(|e| anyhow::anyhow!("{}", e))?;
let target = get_str(args, "target")?; let t = store.resolve_key(get_str(args, "target")?).map_err(|e| anyhow::anyhow!("{}", e))?;
link_add(source, target, prov)? let strength = store.add_link(&s, &t, prov).map_err(|e| anyhow::anyhow!("{}", e))?;
format!("linked {} → {} (strength={:.2})", s, t, strength)
} }
"memory_used" => { "memory_used" => {
let key = get_str(args, "key")?; let key = get_str(args, "key")?;
MemoryNode::mark_used(key) if !store.nodes.contains_key(key) {
.map_err(|e| anyhow::anyhow!("{}", e))? anyhow::bail!("node not found: {}", key);
}
store.mark_used(key);
format!("marked {} as used", key)
} }
"memory_weight_set" => { "memory_weight_set" => {
let key = get_str(args, "key")?; let key = store.resolve_key(get_str(args, "key")?).map_err(|e| anyhow::anyhow!("{}", e))?;
let weight = get_f64(args, "weight")? as f32; let weight = get_f64(args, "weight")? as f32;
weight_set(key, weight)? let (old, new) = store.set_weight(&key, weight).map_err(|e| anyhow::anyhow!("{}", e))?;
format!("weight {} {:.2} → {:.2}", key, old, new)
} }
"memory_supersede" => supersede(args, prov)?, _ => unreachable!(),
_ => anyhow::bail!("Unknown memory tool: {}", name),
}; };
Ok(result)
}
fn link_set(source: &str, target: &str, strength: f32) -> Result<String> {
let mut store = Store::load().map_err(|e| anyhow::anyhow!("{}", e))?;
let source = store.resolve_key(source).map_err(|e| anyhow::anyhow!("{}", e))?;
let target = store.resolve_key(target).map_err(|e| anyhow::anyhow!("{}", e))?;
let old = store.set_link_strength(&source, &target, strength)
.map_err(|e| anyhow::anyhow!("{}", e))?;
store.save().map_err(|e| anyhow::anyhow!("{}", e))?; store.save().map_err(|e| anyhow::anyhow!("{}", e))?;
Ok(format!("set {} ↔ {} strength {:.2} → {:.2}", source, target, old, strength)) Ok(msg)
} }
fn link_add(source: &str, target: &str, prov: &str) -> Result<String> {
let mut store = Store::load().map_err(|e| anyhow::anyhow!("{}", e))?;
let source = store.resolve_key(source).map_err(|e| anyhow::anyhow!("{}", e))?;
let target = store.resolve_key(target).map_err(|e| anyhow::anyhow!("{}", e))?;
let strength = store.add_link(&source, &target, prov)
.map_err(|e| anyhow::anyhow!("{}", e))?;
store.save().map_err(|e| anyhow::anyhow!("{}", e))?;
Ok(format!("linked {} → {} (strength={:.2})", source, target, strength))
}
fn weight_set(key: &str, weight: f32) -> Result<String> {
let mut store = Store::load().map_err(|e| anyhow::anyhow!("{}", e))?;
let resolved = store.resolve_key(key).map_err(|e| anyhow::anyhow!("{}", e))?;
let (old, new) = store.set_weight(&resolved, weight)
.map_err(|e| anyhow::anyhow!("{}", e))?;
store.save().map_err(|e| anyhow::anyhow!("{}", e))?;
Ok(format!("weight {} {:.2} → {:.2}", resolved, old, new))
}
fn supersede(args: &serde_json::Value, prov: &str) -> Result<String> {
let old_key = get_str(args, "old_key")?;
let new_key = get_str(args, "new_key")?;
let reason = args.get("reason").and_then(|v| v.as_str()).unwrap_or("superseded");
// Load old node
let old = MemoryNode::load(old_key)
.ok_or_else(|| anyhow::anyhow!("node not found: {}", old_key))?;
// Prepend superseded notice (strip link footer from content)
let content_only = old.content.split("\n\n---\nLinks:").next().unwrap_or(&old.content);
let notice = format!(
"**SUPERSEDED** by `{}` — {}\n\nOriginal content preserved below for reference.\n\n---\n\n{}",
new_key, reason, content_only.trim()
);
// Write back + set weight
MemoryNode::write(old_key, &notice, Some(prov))
.map_err(|e| anyhow::anyhow!("{}", e))?;
weight_set(old_key, 0.01)?;
Ok(format!("superseded {} → {} ({})", old_key, new_key, reason))
}
/// Helper: get required string argument.
fn get_str<'a>(args: &'a serde_json::Value, name: &'a str) -> Result<&'a str> { fn get_str<'a>(args: &'a serde_json::Value, name: &'a str) -> Result<&'a str> {
args.get(name) args.get(name).and_then(|v| v.as_str()).context(format!("{} is required", name))
.and_then(|v| v.as_str())
.context(format!("{} is required", name))
} }
/// Helper: get required f64 argument.
fn get_f64(args: &serde_json::Value, name: &str) -> Result<f64> { fn get_f64(args: &serde_json::Value, name: &str) -> Result<f64> {
args.get(name) args.get(name).and_then(|v| v.as_f64()).context(format!("{} is required", name))
.and_then(|v| v.as_f64())
.context(format!("{} is required", name))
} }

90
src/hippocampus/memory.rs
View file

@ -2,7 +2,7 @@
// //
// MemoryNode is a lightweight representation of a loaded node: // MemoryNode is a lightweight representation of a loaded node:
// key, content, links, version, weight. Used by the agent for // key, content, links, version, weight. Used by the agent for
// context tracking and by tools for direct store access. // context tracking and by the CLI for rendering.
use super::store::Store; use super::store::Store;
@ -11,24 +11,13 @@ use super::store::Store;
pub struct MemoryNode { pub struct MemoryNode {
pub key: String, pub key: String,
pub content: String, pub content: String,
pub links: Vec<Link>, pub links: Vec<(String, f32)>, // (target_key, strength)
/// Version from the store — used for change detection.
pub version: u32, pub version: u32,
/// Weight in the graph.
pub weight: f32, pub weight: f32,
} }
/// A link to a neighbor node.
#[derive(Debug, Clone)]
pub struct Link {
pub target: String,
pub strength: f32,
/// Whether this link target is already referenced inline in the content.
pub inline: bool,
}
impl MemoryNode { impl MemoryNode {
/// Load a node from the store by key. Returns None if not found. /// Load a node from the store by key.
pub fn load(key: &str) -> Option<Self> { pub fn load(key: &str) -> Option<Self> {
let store = Store::load().ok()?; let store = Store::load().ok()?;
Self::from_store(&store, key) Self::from_store(&store, key)
@ -38,7 +27,6 @@ impl MemoryNode {
pub fn from_store(store: &Store, key: &str) -> Option<Self> { pub fn from_store(store: &Store, key: &str) -> Option<Self> {
let node = store.nodes.get(key)?; let node = store.nodes.get(key)?;
// Collect neighbor strengths
let mut neighbors: std::collections::HashMap<&str, f32> = std::collections::HashMap::new(); let mut neighbors: std::collections::HashMap<&str, f32> = std::collections::HashMap::new();
for r in &store.relations { for r in &store.relations {
if r.deleted { continue; } if r.deleted { continue; }
@ -51,14 +39,10 @@ impl MemoryNode {
} }
} }
let mut links: Vec<Link> = neighbors.into_iter() let mut links: Vec<(String, f32)> = neighbors.into_iter()
.map(|(target, strength)| Link { .map(|(k, s)| (k.to_string(), s))
inline: node.content.contains(target),
target: target.to_string(),
strength,
})
.collect(); .collect();
links.sort_by(|a, b| b.strength.total_cmp(&a.strength)); links.sort_by(|a, b| b.1.total_cmp(&a.1));
Some(MemoryNode { Some(MemoryNode {
key: key.to_string(), key: key.to_string(),
@ -74,16 +58,15 @@ impl MemoryNode {
let mut out = self.content.clone(); let mut out = self.content.clone();
// Footer: links not already referenced inline // Footer: links not already referenced inline
let footer_links: Vec<&Link> = self.links.iter() let footer: Vec<&(String, f32)> = self.links.iter()
.filter(|l| !l.inline) .filter(|(target, _)| !self.content.contains(target.as_str()))
.collect(); .collect();
if !footer_links.is_empty() { if !footer.is_empty() {
let total = footer_links.len(); let total = footer.len();
out.push_str("\n\n---\nLinks:"); out.push_str("\n\n---\nLinks:");
for link in footer_links.iter().take(15) { for (target, strength) in footer.iter().take(15) {
out.push_str(&format!("\n ({:.2}) `poc-memory render {}`", out.push_str(&format!("\n ({:.2}) `poc-memory render {}`", strength, target));
link.strength, link.target));
} }
if total > 15 { if total > 15 {
out.push_str(&format!("\n ... and {} more (`poc-memory graph link {}`)", out.push_str(&format!("\n ... and {} more (`poc-memory graph link {}`)",
@ -92,53 +75,4 @@ impl MemoryNode {
} }
out out
} }
/// Write content to the store and return an updated MemoryNode.
pub fn write(key: &str, content: &str, provenance: Option<&str>) -> Result<Self, String> {
let prov = provenance.unwrap_or("manual");
let mut store = Store::load()?;
store.upsert_provenance(key, content, prov)?;
store.save()?;
Self::from_store(&store, key)
.ok_or_else(|| format!("wrote {} but failed to load back", key))
}
/// Search for nodes matching a query. Returns lightweight results.
pub fn search(query: &str) -> Result<Vec<SearchResult>, String> {
let store = Store::load()?;
let results = super::query::engine::search(query, &store);
Ok(results.into_iter().take(20).map(|hit| SearchResult {
key: hit.key.clone(),
score: hit.activation as f32,
snippet: hit.snippet.unwrap_or_default(),
}).collect())
}
/// Mark a node as used (boosts weight).
pub fn mark_used(key: &str) -> Result<String, String> {
let mut store = Store::load()?;
if !store.nodes.contains_key(key) {
return Err(format!("node not found: {}", key));
}
store.mark_used(key);
store.save()?;
Ok(format!("marked {} as used", key))
}
}
/// A search result — lightweight, not a full node load.
#[derive(Debug, Clone)]
pub struct SearchResult {
pub key: String,
pub score: f32,
pub snippet: String,
}
impl SearchResult {
/// Format for display.
pub fn render(&self) -> String {
format!("({:.2}) {} — {}", self.score, self.key, self.snippet)
}
} }