tools/memory: direct store calls instead of spawning poc-memory
Every memory tool call was spawning a poc-memory subprocess. Now uses
MemoryNode and direct Store API calls:
- memory_render: MemoryNode::load() + render()
- memory_write: MemoryNode::write() via store.upsert_provenance()
- memory_search: MemoryNode::search() via search engine
- memory_links: MemoryNode::load() + iterate links
- memory_link_add: store.add_relation() with Jaccard strength
- memory_link_set: direct relation mutation
- memory_used: store.mark_used()
- memory_weight_set: direct node.weight mutation
- memory_supersede: MemoryNode::load() + write() + weight_set()
No more Command::new("poc-memory") in this module.
Co-Authored-By: Proof of Concept <poc@bcachefs.org>
This commit is contained in:
ProofOfConcept
parent
4cc4952234
commit
2c61a3575d
1 changed files with 130 additions and 61 deletions
|
|
@ -1,15 +1,15 @@
|
|||
// tools/memory.rs — Native memory graph operations
|
||||
//
|
||||
// Structured tool calls for the memory graph, replacing bash
|
||||
// poc-memory commands. Cleaner for LLMs — no shell quoting,
|
||||
// multi-line content as JSON strings, typed parameters.
|
||||
// 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 serde_json::json;
|
||||
use std::io::Write;
|
||||
use std::process::{Command, Stdio};
|
||||
|
||||
use crate::agent::memory::MemoryNode;
|
||||
use crate::agent::types::ToolDef;
|
||||
use crate::store::{self, Store};
|
||||
|
||||
pub fn definitions() -> Vec<ToolDef> {
|
||||
vec![
|
||||
|
|
@ -63,7 +63,7 @@ pub fn definitions() -> Vec<ToolDef> {
|
|||
),
|
||||
ToolDef::new(
|
||||
"memory_links",
|
||||
"Show a node's neighbors with link strengths and clustering coefficients.",
|
||||
"Show a node's neighbors with link strengths.",
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
|
|
@ -176,110 +176,179 @@ pub fn definitions() -> Vec<ToolDef> {
|
|||
]
|
||||
}
|
||||
|
||||
/// Dispatch a memory tool call. Shells out to poc-memory CLI.
|
||||
/// 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");
|
||||
let result = match name {
|
||||
"memory_render" => {
|
||||
let key = get_str(args, "key")?;
|
||||
cmd(&["render", key], provenance)?
|
||||
let node = MemoryNode::load(key)
|
||||
.ok_or_else(|| anyhow::anyhow!("node not found: {}", key))?;
|
||||
node.render()
|
||||
}
|
||||
"memory_write" => {
|
||||
let key = get_str(args, "key")?;
|
||||
let content = get_str(args, "content")?;
|
||||
write_node(key, content, provenance)?
|
||||
let node = MemoryNode::write(key, content, Some(prov))
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
format!("wrote '{}' (v{})", node.key, node.version)
|
||||
}
|
||||
"memory_search" => {
|
||||
let query = get_str(args, "query")?;
|
||||
cmd(&["search", query], provenance)?
|
||||
let results = MemoryNode::search(query)
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
if results.is_empty() {
|
||||
"no results".to_string()
|
||||
} else {
|
||||
results.iter().map(|r| r.render()).collect::<Vec<_>>().join("\n")
|
||||
}
|
||||
}
|
||||
"memory_links" => {
|
||||
let key = get_str(args, "key")?;
|
||||
cmd(&["graph", "link", key], provenance)?
|
||||
let node = MemoryNode::load(key)
|
||||
.ok_or_else(|| anyhow::anyhow!("node not found: {}", key))?;
|
||||
let mut out = format!("Neighbors of '{}':\n", key);
|
||||
for link in &node.links {
|
||||
out.push_str(&format!(" ({:.2}) {}{}\n",
|
||||
link.strength, link.target,
|
||||
if link.inline { " [inline]" } else { "" }));
|
||||
}
|
||||
out
|
||||
}
|
||||
"memory_link_set" => {
|
||||
let source = get_str(args, "source")?;
|
||||
let target = get_str(args, "target")?;
|
||||
let strength = get_f64(args, "strength")?;
|
||||
cmd(&["graph", "link-set", source, target, &format!("{:.2}", strength)], provenance)?
|
||||
let strength = get_f64(args, "strength")? as f32;
|
||||
link_set(source, target, strength)?
|
||||
}
|
||||
"memory_link_add" => {
|
||||
let source = get_str(args, "source")?;
|
||||
let target = get_str(args, "target")?;
|
||||
cmd(&["graph", "link-add", source, target], provenance)?
|
||||
link_add(source, target, prov)?
|
||||
}
|
||||
"memory_used" => {
|
||||
let key = get_str(args, "key")?;
|
||||
cmd(&["used", key], provenance)?
|
||||
MemoryNode::mark_used(key)
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?
|
||||
}
|
||||
"memory_weight_set" => {
|
||||
let key = get_str(args, "key")?;
|
||||
let weight = get_f64(args, "weight")?;
|
||||
cmd(&["weight-set", key, &format!("{:.2}", weight)], provenance)?
|
||||
let weight = get_f64(args, "weight")? as f32;
|
||||
weight_set(key, weight)?
|
||||
}
|
||||
"memory_supersede" => supersede(args, provenance)?,
|
||||
"memory_supersede" => supersede(args, prov)?,
|
||||
_ => anyhow::bail!("Unknown memory tool: {}", name),
|
||||
};
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// Run poc-memory command and return stdout.
|
||||
fn cmd(args: &[&str], provenance: Option<&str>) -> Result<String> {
|
||||
let mut cmd = Command::new("poc-memory");
|
||||
cmd.args(args);
|
||||
if let Some(prov) = provenance {
|
||||
cmd.env("POC_PROVENANCE", prov);
|
||||
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 strength = strength.clamp(0.01, 1.0);
|
||||
|
||||
let mut found = false;
|
||||
let mut first = true;
|
||||
for rel in &mut store.relations {
|
||||
if rel.deleted { continue; }
|
||||
if (rel.source_key == source && rel.target_key == target)
|
||||
|| (rel.source_key == target && rel.target_key == source)
|
||||
{
|
||||
if first {
|
||||
let old = rel.strength;
|
||||
rel.strength = strength;
|
||||
first = false;
|
||||
found = true;
|
||||
if (old - strength).abs() < 0.001 {
|
||||
return Ok(format!("{} ↔ {} already at {:.2}", source, target, strength));
|
||||
}
|
||||
} else {
|
||||
rel.deleted = true; // deduplicate
|
||||
}
|
||||
}
|
||||
}
|
||||
let output = cmd.output().context("run poc-memory")?;
|
||||
let stdout = String::from_utf8_lossy(&output.stdout);
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
if output.status.success() {
|
||||
Ok(stdout.to_string())
|
||||
} else {
|
||||
Ok(format!("{}{}", stdout, stderr))
|
||||
|
||||
if !found {
|
||||
anyhow::bail!("no link found between {} and {}", source, target);
|
||||
}
|
||||
|
||||
store.save().map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
Ok(format!("set {} ↔ {} strength to {:.2}", source, target, strength))
|
||||
}
|
||||
|
||||
/// Write content to a node via stdin.
|
||||
fn write_node(key: &str, content: &str, provenance: Option<&str>) -> Result<String> {
|
||||
let mut cmd = Command::new("poc-memory");
|
||||
cmd.args(["write", key])
|
||||
.stdin(Stdio::piped())
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped());
|
||||
if let Some(prov) = provenance {
|
||||
cmd.env("POC_PROVENANCE", prov);
|
||||
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))?;
|
||||
|
||||
// Check for existing link
|
||||
let exists = store.relations.iter().any(|r|
|
||||
!r.deleted &&
|
||||
((r.source_key == source && r.target_key == target) ||
|
||||
(r.source_key == target && r.target_key == source)));
|
||||
if exists {
|
||||
return Ok(format!("link already exists: {} ↔ {}", source, target));
|
||||
}
|
||||
let mut child = cmd.spawn().context("spawn poc-memory write")?;
|
||||
child.stdin.take().unwrap().write_all(content.as_bytes())
|
||||
.context("write content to stdin")?;
|
||||
let output = child.wait_with_output().context("wait poc-memory write")?;
|
||||
Ok(String::from_utf8_lossy(&output.stdout).to_string()
|
||||
+ &String::from_utf8_lossy(&output.stderr))
|
||||
|
||||
let source_uuid = store.nodes.get(&source)
|
||||
.map(|n| n.uuid)
|
||||
.ok_or_else(|| anyhow::anyhow!("source not found: {}", source))?;
|
||||
let target_uuid = store.nodes.get(&target)
|
||||
.map(|n| n.uuid)
|
||||
.ok_or_else(|| anyhow::anyhow!("target not found: {}", target))?;
|
||||
|
||||
// Compute initial strength from Jaccard similarity
|
||||
let graph = store.build_graph();
|
||||
let jaccard = graph.jaccard(&source, &target);
|
||||
let strength = (jaccard * 3.0).clamp(0.1, 1.0) as f32;
|
||||
|
||||
let mut rel = store::new_relation(
|
||||
source_uuid, target_uuid,
|
||||
store::RelationType::Link, strength,
|
||||
&source, &target,
|
||||
);
|
||||
rel.provenance = prov.to_string();
|
||||
store.add_relation(rel).map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
store.save().map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
Ok(format!("linked {} → {} (strength={:.2})", source, target, strength))
|
||||
}
|
||||
|
||||
/// Handle memory_supersede - reads old node, prepends notice, writes back, sets weight.
|
||||
fn supersede(args: &serde_json::Value, provenance: Option<&str>) -> Result<String> {
|
||||
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 weight = weight.clamp(0.01, 1.0);
|
||||
|
||||
let node = store.nodes.get_mut(&resolved)
|
||||
.ok_or_else(|| anyhow::anyhow!("node not found: {}", resolved))?;
|
||||
let old = node.weight;
|
||||
node.weight = weight;
|
||||
store.save().map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
Ok(format!("weight {} {:.2} → {:.2}", resolved, old, weight))
|
||||
}
|
||||
|
||||
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");
|
||||
|
||||
// Read old node
|
||||
let old_content = cmd(&["render", old_key], provenance)?;
|
||||
let content_only = old_content.split("\n\n---\nLinks:").next().unwrap_or(&old_content);
|
||||
// Load old node
|
||||
let old = MemoryNode::load(old_key)
|
||||
.ok_or_else(|| anyhow::anyhow!("node not found: {}", old_key))?;
|
||||
|
||||
// Prepend superseded notice
|
||||
// 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
|
||||
let write_result = write_node(old_key, ¬ice, provenance)?;
|
||||
// Write back + set weight
|
||||
MemoryNode::write(old_key, ¬ice, Some(prov))
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
weight_set(old_key, 0.01)?;
|
||||
|
||||
// Set weight to 0.01
|
||||
let weight_result = cmd(&["weight-set", old_key, "0.01"], provenance)?;
|
||||
|
||||
Ok(format!("{}\n{}", write_result.trim(), weight_result.trim()))
|
||||
Ok(format!("superseded {} → {} ({})", old_key, new_key, reason))
|
||||
}
|
||||
|
||||
/// Helper: get required string argument.
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue