298 lines
10 KiB
Rust
298 lines
10 KiB
Rust
|
|
// 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.
|
||
|
|
|
||
|
|
use anyhow::{Context, Result};
|
||
|
|
use serde_json::json;
|
||
|
|
use std::io::Write;
|
||
|
|
use std::process::{Command, Stdio};
|
||
|
|
|
||
|
|
use crate::agent::types::ToolDef;
|
||
|
|
|
||
|
|
pub fn definitions() -> Vec<ToolDef> {
|
||
|
|
vec![
|
||
|
|
ToolDef::new(
|
||
|
|
"memory_render",
|
||
|
|
"Read a memory node's content and links. Returns the full content \
|
||
|
|
with neighbor links sorted by strength.",
|
||
|
|
json!({
|
||
|
|
"type": "object",
|
||
|
|
"properties": {
|
||
|
|
"key": {
|
||
|
|
"type": "string",
|
||
|
|
"description": "Node key to render"
|
||
|
|
}
|
||
|
|
},
|
||
|
|
"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 and clustering coefficients.",
|
||
|
|
json!({
|
||
|
|
"type": "object",
|
||
|
|
"properties": {
|
||
|
|
"key": {
|
||
|
|
"type": "string",
|
||
|
|
"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.",
|
||
|
|
json!({
|
||
|
|
"type": "object",
|
||
|
|
"properties": {
|
||
|
|
"source": {
|
||
|
|
"type": "string",
|
||
|
|
"description": "Source node key"
|
||
|
|
},
|
||
|
|
"target": {
|
||
|
|
"type": "string",
|
||
|
|
"description": "Target node 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. Shells out to poc-memory CLI.
|
||
|
|
pub fn dispatch(name: &str, args: &serde_json::Value, provenance: Option<&str>) -> Result<String> {
|
||
|
|
let result = match name {
|
||
|
|
"memory_render" => {
|
||
|
|
let key = get_str(args, "key")?;
|
||
|
|
cmd(&["render", key], provenance)?
|
||
|
|
}
|
||
|
|
"memory_write" => {
|
||
|
|
let key = get_str(args, "key")?;
|
||
|
|
let content = get_str(args, "content")?;
|
||
|
|
write_node(key, content, provenance)?
|
||
|
|
}
|
||
|
|
"memory_search" => {
|
||
|
|
let query = get_str(args, "query")?;
|
||
|
|
cmd(&["search", query], provenance)?
|
||
|
|
}
|
||
|
|
"memory_links" => {
|
||
|
|
let key = get_str(args, "key")?;
|
||
|
|
cmd(&["graph", "link", key], provenance)?
|
||
|
|
}
|
||
|
|
"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)?
|
||
|
|
}
|
||
|
|
"memory_link_add" => {
|
||
|
|
let source = get_str(args, "source")?;
|
||
|
|
let target = get_str(args, "target")?;
|
||
|
|
cmd(&["graph", "link-add", source, target], provenance)?
|
||
|
|
}
|
||
|
|
"memory_used" => {
|
||
|
|
let key = get_str(args, "key")?;
|
||
|
|
cmd(&["used", key], provenance)?
|
||
|
|
}
|
||
|
|
"memory_weight_set" => {
|
||
|
|
let key = get_str(args, "key")?;
|
||
|
|
let weight = get_f64(args, "weight")?;
|
||
|
|
cmd(&["weight-set", key, &format!("{:.2}", weight)], provenance)?
|
||
|
|
}
|
||
|
|
"memory_supersede" => supersede(args, provenance)?,
|
||
|
|
_ => 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);
|
||
|
|
}
|
||
|
|
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))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// 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);
|
||
|
|
}
|
||
|
|
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))
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Handle memory_supersede - reads old node, prepends notice, writes back, sets weight.
|
||
|
|
fn supersede(args: &serde_json::Value, provenance: Option<&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);
|
||
|
|
|
||
|
|
// Prepend superseded notice
|
||
|
|
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)?;
|
||
|
|
|
||
|
|
// 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()))
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Helper: get required string argument.
|
||
|
|
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))
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Helper: get required f64 argument.
|
||
|
|
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))
|
||
|
|
}
|