Skip to main content

walrus_memory/
tool.rs

1//! Input parameters for the memory tools.
2
3use schemars::JsonSchema;
4use serde::Deserialize;
5use wcore::{
6    agent::{AsTool, ToolDescription},
7    protocol::ext::ToolDef,
8};
9
10#[derive(Deserialize, JsonSchema)]
11pub struct Recall {
12    /// Batch of search queries to run against memory.
13    pub queries: Vec<String>,
14    /// Maximum number of results per query (default: 5).
15    pub limit: Option<u32>,
16}
17
18impl ToolDescription for Recall {
19    const DESCRIPTION: &'static str =
20        "Search memory by one or more queries. Returns relevant entities and graph connections.";
21}
22
23#[derive(Deserialize, JsonSchema)]
24pub struct Extract {
25    /// Entities to upsert.
26    pub entities: Vec<ExtractEntity>,
27    /// Relations to upsert.
28    pub relations: Vec<ExtractRelation>,
29}
30
31#[derive(Deserialize, JsonSchema)]
32pub struct ExtractEntity {
33    /// Human-readable key/name for the entity.
34    pub key: String,
35    /// Value/content to store.
36    pub value: String,
37    /// Optional entity type (e.g. "fact", "person", "preference").
38    pub entity_type: Option<String>,
39}
40
41#[derive(Deserialize, JsonSchema)]
42pub struct ExtractRelation {
43    /// Key of the source entity.
44    pub source: String,
45    /// Relation label (e.g. "knows", "prefers", "related_to").
46    pub relation: String,
47    /// Key of the target entity.
48    pub target: String,
49}
50
51impl ToolDescription for Extract {
52    const DESCRIPTION: &'static str =
53        "Batch upsert entities and relations into memory. Internal tool for extraction.";
54}
55
56/// Build agent-visible tool defs: only `recall`.
57pub fn tool_defs() -> Vec<ToolDef> {
58    let t = Recall::as_tool();
59    vec![ToolDef {
60        name: t.name.to_string(),
61        description: t.description.to_string(),
62        parameters: serde_json::to_vec(&t.parameters).expect("schema serialization"),
63        strict: t.strict,
64    }]
65}
66
67/// Build all tool defs including internal ones (for extension RegisterTools response).
68///
69/// The daemon needs `extract` in the ToolSchemas so `infer_fulfill` can
70/// provide it to the extraction LLM.
71pub fn all_tool_defs() -> Vec<ToolDef> {
72    [Recall::as_tool(), Extract::as_tool()]
73        .into_iter()
74        .map(|t| ToolDef {
75            name: t.name.to_string(),
76            description: t.description.to_string(),
77            parameters: serde_json::to_vec(&t.parameters).expect("schema serialization"),
78            strict: t.strict,
79        })
80        .collect()
81}