consciousness/src/agent/tools/memory.rs

260 lines
13 KiB
Rust
Raw Normal View History

// tools/memory.rs — Native memory graph operations
//
// Direct library calls into the store — no subprocess spawning.
use anyhow::{Context, Result};
use serde_json::json;
use crate::hippocampus::memory::MemoryNode;
use crate::agent::types::ToolDef;
use crate::store::Store;
pub fn definitions() -> Vec<ToolDef> {
vec![
ToolDef::new("memory_render",
"Read a memory node's content and links.",
json!({"type":"object","properties":{"key":{"type":"string","description":"Node key"}},"required":["key"]})),
ToolDef::new("memory_write",
"Create or update a memory node.",
json!({"type":"object","properties":{"key":{"type":"string","description":"Node key"},"content":{"type":"string","description":"Full content (markdown)"}},"required":["key","content"]})),
ToolDef::new("memory_search",
"Search the memory graph 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.",
json!({"type":"object","properties":{"key":{"type":"string","description":"Node key"}},"required":["key"]})),
ToolDef::new("memory_link_set",
"Set link strength between two nodes.",
json!({"type":"object","properties":{"source":{"type":"string"},"target":{"type":"string"},"strength":{"type":"number","description":"0.01 to 1.0"}},"required":["source","target","strength"]})),
ToolDef::new("memory_link_add",
"Add a new link between two nodes.",
json!({"type":"object","properties":{"source":{"type":"string"},"target":{"type":"string"}},"required":["source","target"]})),
ToolDef::new("memory_used",
"Mark a node as useful (boosts weight).",
json!({"type":"object","properties":{"key":{"type":"string","description":"Node key"}},"required":["key"]})),
ToolDef::new("memory_weight_set",
"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"]})),
ToolDef::new("memory_supersede",
"Mark a node as superseded by another (sets weight to 0.01).",
json!({"type":"object","properties":{"old_key":{"type":"string"},"new_key":{"type":"string"},"reason":{"type":"string"}},"required":["old_key","new_key"]})),
ToolDef::new("memory_query",
"Run a structured query against the memory graph. Supports filtering, \
sorting, field selection. Examples: \"degree > 10 | sort weight | limit 5\", \
\"neighbors('identity') | select strength\", \"key ~ 'journal.*' | count\"",
json!({"type":"object","properties":{"query":{"type":"string","description":"Query expression"}},"required":["query"]})),
ToolDef::new("output",
"Produce a named output value. Use this to pass structured results \
between steps subsequent prompts can see these in the conversation history.",
json!({"type":"object","properties":{
"key":{"type":"string","description":"Output name (e.g. 'relevant_memories')"},
"value":{"type":"string","description":"Output value"}
},"required":["key","value"]})),
ToolDef::new("journal_tail",
"Read the last N journal entries (default 1).",
json!({"type":"object","properties":{
"count":{"type":"integer","description":"Number of entries (default 1)"}
}})),
ToolDef::new("journal_new",
"Start a new journal entry with a ## heading and body.",
json!({"type":"object","properties":{
"title":{"type":"string","description":"Entry title (becomes ## YYYY-MM-DDTHH:MM — title)"},
"body":{"type":"string","description":"Entry body (2-3 paragraphs)"}
},"required":["title","body"]})),
ToolDef::new("journal_update",
"Append text to the most recent journal entry (same thread continuing).",
json!({"type":"object","properties":{
"body":{"type":"string","description":"Text to append to the last entry"}
},"required":["body"]})),
]
}
/// Dispatch a memory tool call. Direct library calls, no subprocesses.
pub fn dispatch(name: &str, args: &serde_json::Value, provenance: Option<&str>) -> Result<String> {
let prov = provenance.unwrap_or("manual");
match name {
"memory_render" => {
let key = get_str(args, "key")?;
Ok(MemoryNode::load(key)
.ok_or_else(|| anyhow::anyhow!("node not found: {}", key))?
.render())
}
"memory_write" => {
let key = get_str(args, "key")?;
let content = get_str(args, "content")?;
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))?;
store.save().map_err(|e| anyhow::anyhow!("{}", e))?;
Ok(format!("{} '{}'", result, key))
}
"memory_search" => {
let query = get_str(args, "query")?;
let store = Store::load().map_err(|e| anyhow::anyhow!("{}", e))?;
let results = crate::search::search(query, &store);
if results.is_empty() {
Ok("no results".into())
} else {
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" => {
let key = get_str(args, "key")?;
let node = MemoryNode::load(key)
.ok_or_else(|| anyhow::anyhow!("node not found: {}", key))?;
let mut out = format!("Neighbors of '{}':\n", key);
for (target, strength) in &node.links {
out.push_str(&format!(" ({:.2}) {}\n", strength, target));
}
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))
}
"memory_query" => {
let query = get_str(args, "query")?;
let store = Store::load().map_err(|e| anyhow::anyhow!("{}", e))?;
let graph = store.build_graph();
crate::query_parser::query_to_string(&store, &graph, query)
.map_err(|e| anyhow::anyhow!("{}", e))
}
"output" => {
let key = get_str(args, "key")?;
if key.starts_with("pid-") || key.contains('/') || key.contains("..") {
anyhow::bail!("invalid output key: {}", key);
}
let value = get_str(args, "value")?;
let dir = std::env::var("POC_AGENT_OUTPUT_DIR")
.map_err(|_| anyhow::anyhow!("no output directory set"))?;
let path = std::path::Path::new(&dir).join(key);
std::fs::write(&path, value)
.with_context(|| format!("writing output {}", path.display()))?;
Ok(format!("{}: {}", key, value))
}
"journal_tail" => {
let count = args.get("count").and_then(|v| v.as_u64()).unwrap_or(1) as usize;
let store = Store::load().map_err(|e| anyhow::anyhow!("{}", e))?;
let content = store.nodes.get("journal")
.map(|n| n.content.as_str())
.unwrap_or("");
let mut entries: Vec<&str> = Vec::new();
let mut remaining = content;
while let Some(pos) = remaining.rfind("\n## ") {
entries.push(&remaining[pos + 1..]);
remaining = &remaining[..pos];
if entries.len() >= count { break; }
}
if entries.len() < count && remaining.starts_with("## ") {
entries.push(remaining);
}
entries.reverse();
if entries.is_empty() {
Ok("(no journal entries)".into())
} else {
Ok(entries.join("\n\n"))
}
}
"journal_new" => {
let title = get_str(args, "title")?;
let body = get_str(args, "body")?;
let ts = chrono::Local::now().format("%Y-%m-%dT%H:%M");
let entry = format!("## {}{}\n\n{}", ts, title, body);
let mut store = Store::load().map_err(|e| anyhow::anyhow!("{}", e))?;
let existing = store.nodes.get("journal")
.map(|n| n.content.clone())
.unwrap_or_default();
let new_content = if existing.is_empty() {
entry.clone()
} else {
format!("{}\n\n{}", existing.trim_end(), entry)
};
store.upsert_provenance("journal", &new_content, prov)
.map_err(|e| anyhow::anyhow!("{}", e))?;
store.save().map_err(|e| anyhow::anyhow!("{}", e))?;
let word_count = body.split_whitespace().count();
Ok(format!("New entry '{}' ({} words)", title, word_count))
}
"journal_update" => {
let body = get_str(args, "body")?;
let mut store = Store::load().map_err(|e| anyhow::anyhow!("{}", e))?;
let existing = store.nodes.get("journal")
.map(|n| n.content.clone())
.unwrap_or_default();
if existing.is_empty() {
anyhow::bail!("no journal entry to update — use journal_new first");
}
let new_content = format!("{}\n\n{}", existing.trim_end(), body);
store.upsert_provenance("journal", &new_content, prov)
.map_err(|e| anyhow::anyhow!("{}", e))?;
store.save().map_err(|e| anyhow::anyhow!("{}", e))?;
let word_count = body.split_whitespace().count();
Ok(format!("Updated last entry (+{} words)", word_count))
}
_ => 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" => {
let s = store.resolve_key(get_str(args, "source")?).map_err(|e| anyhow::anyhow!("{}", e))?;
let t = store.resolve_key(get_str(args, "target")?).map_err(|e| anyhow::anyhow!("{}", e))?;
let strength = get_f64(args, "strength")? as f32;
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" => {
let s = store.resolve_key(get_str(args, "source")?).map_err(|e| anyhow::anyhow!("{}", e))?;
let t = store.resolve_key(get_str(args, "target")?).map_err(|e| anyhow::anyhow!("{}", e))?;
let strength = store.add_link(&s, &t, prov).map_err(|e| anyhow::anyhow!("{}", e))?;
format!("linked {}{} (strength={:.2})", s, t, strength)
}
"memory_used" => {
let key = get_str(args, "key")?;
if !store.nodes.contains_key(key) {
anyhow::bail!("node not found: {}", key);
}
store.mark_used(key);
format!("marked {} as used", key)
}
"memory_weight_set" => {
let key = store.resolve_key(get_str(args, "key")?).map_err(|e| anyhow::anyhow!("{}", e))?;
let weight = get_f64(args, "weight")? as f32;
let (old, new) = store.set_weight(&key, weight).map_err(|e| anyhow::anyhow!("{}", e))?;
format!("weight {} {:.2}{:.2}", key, old, new)
}
_ => unreachable!(),
};
store.save().map_err(|e| anyhow::anyhow!("{}", e))?;
Ok(msg)
}
fn get_str<'a>(args: &'a serde_json::Value, name: &'a str) -> Result<&'a str> {
args.get(name).and_then(|v| v.as_str()).context(format!("{} is required", name))
}
fn get_f64(args: &serde_json::Value, name: &str) -> Result<f64> {
args.get(name).and_then(|v| v.as_f64()).context(format!("{} is required", name))
}