forked from kent/consciousness
memory tools: generate public typed API via macro
The memory_tool! macro now generates two functions:
- jsonargs_*() - internal, takes JSON args for dispatch table
- pub fn name() - typed args, handles RPC-vs-local automatically
Callers can now use typed Rust API:
memory::write(Some(&agent), "key", "content").await?;
memory::query(None, "all | type:semantic", Some("full")).await?;
No more manual JSON construction for memory tool calls.
Co-Authored-By: Proof of Concept <poc@bcachefs.org>
This commit is contained in:
Kent Overstreet
parent
4560ba9230
commit
933221f482
1 changed files with 131 additions and 48 deletions
|
|
@ -75,33 +75,14 @@ fn get_provenance(args: &serde_json::Value) -> String {
|
|||
|
||||
// ── Macro for generating tool wrappers ─────────────────────────
|
||||
//
|
||||
// memory_tool!(name, mut, arg1: str, arg2: f32, arg3: ?str)
|
||||
// memory_tool!(name, mut, arg1: [str], arg2: [Option<bool>])
|
||||
// - mut/ref for store mutability
|
||||
// - type suffixes: str, f32, f64, u64, i64, bool
|
||||
// - ?type for optional args with default
|
||||
// - generates jsonargs_* (internal, JSON args) and public typed API
|
||||
|
||||
macro_rules! memory_tool {
|
||||
// Mutable store variant
|
||||
($name:ident, mut $(, $($arg:ident : [$($typ:tt)+]),* $(,)?)?) => {
|
||||
async fn $name(args: &serde_json::Value) -> Result<String> {
|
||||
$($(let $arg = memory_tool!(@extract args, $arg, $($typ)+);)*)?
|
||||
let prov = get_provenance(args);
|
||||
let arc = cached_store().await?;
|
||||
let mut store = arc.lock().await;
|
||||
crate::hippocampus::$name(&mut store, &prov $($(, $arg)*)?)
|
||||
}
|
||||
};
|
||||
// Immutable store variant
|
||||
($name:ident, ref $(, $($arg:ident : [$($typ:tt)+]),* $(,)?)?) => {
|
||||
async fn $name(args: &serde_json::Value) -> Result<String> {
|
||||
$($(let $arg = memory_tool!(@extract args, $arg, $($typ)+);)*)?
|
||||
let prov = get_provenance(args);
|
||||
let arc = cached_store().await?;
|
||||
let store = arc.lock().await;
|
||||
crate::hippocampus::$name(&store, &prov $($(, $arg)*)?)
|
||||
}
|
||||
};
|
||||
// Required extractors - fail if missing
|
||||
// ── Helper rules (must come first) ─────────────────────────────
|
||||
|
||||
// Extract from JSON
|
||||
(@extract $args:ident, $name:ident, str) => {
|
||||
get_str($args, stringify!($name))?
|
||||
};
|
||||
|
|
@ -114,8 +95,6 @@ macro_rules! memory_tool {
|
|||
.map(|arr| arr.iter().filter_map(|v| v.as_str().map(String::from)).collect::<Vec<_>>())
|
||||
.unwrap_or_default()
|
||||
};
|
||||
|
||||
// Optional extractors - return Option<T>
|
||||
(@extract $args:ident, $name:ident, Option<&str>) => {
|
||||
$args.get(stringify!($name)).and_then(|v| v.as_str())
|
||||
};
|
||||
|
|
@ -137,6 +116,110 @@ macro_rules! memory_tool {
|
|||
(@extract $args:ident, $name:ident, Option<f64>) => {
|
||||
$args.get(stringify!($name)).and_then(|v| v.as_f64())
|
||||
};
|
||||
|
||||
// Parameter types for function signatures
|
||||
(@param_type str) => { &str };
|
||||
(@param_type f32) => { f32 };
|
||||
(@param_type Vec<String>) => { Vec<String> };
|
||||
(@param_type Option<&str>) => { Option<&str> };
|
||||
(@param_type Option<bool>) => { Option<bool> };
|
||||
(@param_type Option<u64>) => { Option<u64> };
|
||||
(@param_type Option<i64>) => { Option<i64> };
|
||||
(@param_type Option<usize>) => { Option<usize> };
|
||||
(@param_type Option<u32>) => { Option<u32> };
|
||||
(@param_type Option<f64>) => { Option<f64> };
|
||||
|
||||
// Serialize to JSON for RPC
|
||||
(@insert_json $map:ident, $name:ident, str) => {
|
||||
$map.insert(stringify!($name).into(), serde_json::json!($name));
|
||||
};
|
||||
(@insert_json $map:ident, $name:ident, f32) => {
|
||||
$map.insert(stringify!($name).into(), serde_json::json!($name));
|
||||
};
|
||||
(@insert_json $map:ident, $name:ident, Vec<String>) => {
|
||||
$map.insert(stringify!($name).into(), serde_json::json!($name));
|
||||
};
|
||||
(@insert_json $map:ident, $name:ident, Option<&str>) => {
|
||||
if let Some(v) = $name { $map.insert(stringify!($name).into(), serde_json::json!(v)); }
|
||||
};
|
||||
(@insert_json $map:ident, $name:ident, Option<bool>) => {
|
||||
if let Some(v) = $name { $map.insert(stringify!($name).into(), serde_json::json!(v)); }
|
||||
};
|
||||
(@insert_json $map:ident, $name:ident, Option<u64>) => {
|
||||
if let Some(v) = $name { $map.insert(stringify!($name).into(), serde_json::json!(v)); }
|
||||
};
|
||||
(@insert_json $map:ident, $name:ident, Option<i64>) => {
|
||||
if let Some(v) = $name { $map.insert(stringify!($name).into(), serde_json::json!(v)); }
|
||||
};
|
||||
(@insert_json $map:ident, $name:ident, Option<usize>) => {
|
||||
if let Some(v) = $name { $map.insert(stringify!($name).into(), serde_json::json!(v)); }
|
||||
};
|
||||
(@insert_json $map:ident, $name:ident, Option<u32>) => {
|
||||
if let Some(v) = $name { $map.insert(stringify!($name).into(), serde_json::json!(v)); }
|
||||
};
|
||||
(@insert_json $map:ident, $name:ident, Option<f64>) => {
|
||||
if let Some(v) = $name { $map.insert(stringify!($name).into(), serde_json::json!(v)); }
|
||||
};
|
||||
|
||||
// ── Main rules ─────────────────────────────────────────────────
|
||||
|
||||
// Mutable store variant
|
||||
($name:ident, mut $(, $($arg:ident : [$($typ:tt)+]),* $(,)?)?) => {
|
||||
paste::paste! {
|
||||
async fn [<jsonargs_ $name>](args: &serde_json::Value) -> Result<String> {
|
||||
$($(let $arg = memory_tool!(@extract args, $arg, $($typ)+);)*)?
|
||||
let prov = get_provenance(args);
|
||||
let arc = cached_store().await?;
|
||||
let mut store = arc.lock().await;
|
||||
crate::hippocampus::$name(&mut store, &prov $($(, $arg)*)?)
|
||||
}
|
||||
|
||||
pub async fn $name(agent: Option<&crate::agent::Agent> $($(, $arg: memory_tool!(@param_type $($typ)+))*)?) -> Result<String> {
|
||||
if !is_daemon() {
|
||||
#[allow(unused_mut)]
|
||||
let mut map = serde_json::Map::new();
|
||||
$($(memory_tool!(@insert_json map, $arg, $($typ)+);)*)?
|
||||
return crate::mcp_server::memory_rpc(concat!("memory_", stringify!($name)), serde_json::Value::Object(map));
|
||||
}
|
||||
let prov = match agent {
|
||||
Some(a) => a.state.lock().await.provenance.clone(),
|
||||
None => "manual".to_string(),
|
||||
};
|
||||
let arc = cached_store().await?;
|
||||
let mut store = arc.lock().await;
|
||||
crate::hippocampus::$name(&mut store, &prov $($(, $arg)*)?)
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// Immutable store variant
|
||||
($name:ident, ref $(, $($arg:ident : [$($typ:tt)+]),* $(,)?)?) => {
|
||||
paste::paste! {
|
||||
async fn [<jsonargs_ $name>](args: &serde_json::Value) -> Result<String> {
|
||||
$($(let $arg = memory_tool!(@extract args, $arg, $($typ)+);)*)?
|
||||
let prov = get_provenance(args);
|
||||
let arc = cached_store().await?;
|
||||
let store = arc.lock().await;
|
||||
crate::hippocampus::$name(&store, &prov $($(, $arg)*)?)
|
||||
}
|
||||
|
||||
pub async fn $name(agent: Option<&crate::agent::Agent> $($(, $arg: memory_tool!(@param_type $($typ)+))*)?) -> Result<String> {
|
||||
if !is_daemon() {
|
||||
#[allow(unused_mut)]
|
||||
let mut map = serde_json::Map::new();
|
||||
$($(memory_tool!(@insert_json map, $arg, $($typ)+);)*)?
|
||||
return crate::mcp_server::memory_rpc(concat!("memory_", stringify!($name)), serde_json::Value::Object(map));
|
||||
}
|
||||
let prov = match agent {
|
||||
Some(a) => a.state.lock().await.provenance.clone(),
|
||||
None => "manual".to_string(),
|
||||
};
|
||||
let arc = cached_store().await?;
|
||||
let store = arc.lock().await;
|
||||
crate::hippocampus::$name(&store, &prov $($(, $arg)*)?)
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
// ── Memory tools ───────────────────────────────────────────────
|
||||
|
|
@ -193,28 +276,28 @@ async fn dispatch(
|
|||
|
||||
// Daemon path - dispatch to implementation
|
||||
match tool_name {
|
||||
"memory_render" => render(&args).await,
|
||||
"memory_write" => write(&args).await,
|
||||
"memory_search" => search(&args).await,
|
||||
"memory_links" => links(&args).await,
|
||||
"memory_link_set" => link_set(&args).await,
|
||||
"memory_link_add" => link_add(&args).await,
|
||||
"memory_delete" => delete(&args).await,
|
||||
"memory_history" => history(&args).await,
|
||||
"memory_weight_set" => weight_set(&args).await,
|
||||
"memory_rename" => rename(&args).await,
|
||||
"memory_supersede" => supersede(&args).await,
|
||||
"memory_query" => query(&args).await,
|
||||
"graph_topology" => graph_topology(&args).await,
|
||||
"graph_health" => graph_health(&args).await,
|
||||
"graph_communities" => graph_communities(&args).await,
|
||||
"graph_normalize_strengths" => graph_normalize_strengths(&args).await,
|
||||
"graph_trace" => graph_trace(&args).await,
|
||||
"graph_link_impact" => graph_link_impact(&args).await,
|
||||
"graph_hubs" => graph_hubs(&args).await,
|
||||
"journal_tail" => journal_tail(&args).await,
|
||||
"journal_new" => journal_new(&args).await,
|
||||
"journal_update" => journal_update(&args).await,
|
||||
"memory_render" => jsonargs_render(&args).await,
|
||||
"memory_write" => jsonargs_write(&args).await,
|
||||
"memory_search" => jsonargs_search(&args).await,
|
||||
"memory_links" => jsonargs_links(&args).await,
|
||||
"memory_link_set" => jsonargs_link_set(&args).await,
|
||||
"memory_link_add" => jsonargs_link_add(&args).await,
|
||||
"memory_delete" => jsonargs_delete(&args).await,
|
||||
"memory_history" => jsonargs_history(&args).await,
|
||||
"memory_weight_set" => jsonargs_weight_set(&args).await,
|
||||
"memory_rename" => jsonargs_rename(&args).await,
|
||||
"memory_supersede" => jsonargs_supersede(&args).await,
|
||||
"memory_query" => jsonargs_query(&args).await,
|
||||
"graph_topology" => jsonargs_graph_topology(&args).await,
|
||||
"graph_health" => jsonargs_graph_health(&args).await,
|
||||
"graph_communities" => jsonargs_graph_communities(&args).await,
|
||||
"graph_normalize_strengths" => jsonargs_graph_normalize_strengths(&args).await,
|
||||
"graph_trace" => jsonargs_graph_trace(&args).await,
|
||||
"graph_link_impact" => jsonargs_graph_link_impact(&args).await,
|
||||
"graph_hubs" => jsonargs_graph_hubs(&args).await,
|
||||
"journal_tail" => jsonargs_journal_tail(&args).await,
|
||||
"journal_new" => jsonargs_journal_new(&args).await,
|
||||
"journal_update" => jsonargs_journal_update(&args).await,
|
||||
_ => anyhow::bail!("unknown tool: {}", tool_name),
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue