Skip to main content

agentic_memory_mcp/tools/
memory_add.rs

1//! Tool: memory_add — Add a cognitive event to the memory graph.
2
3use std::sync::Arc;
4use tokio::sync::Mutex;
5
6use serde::Deserialize;
7use serde_json::{json, Value};
8
9use agentic_memory::{EdgeType, EventType};
10
11use crate::session::SessionManager;
12use crate::types::{McpError, McpResult, ToolCallResult, ToolDefinition};
13
14/// Input parameters for memory_add.
15#[derive(Debug, Deserialize)]
16struct AddParams {
17    event_type: String,
18    content: String,
19    #[serde(default = "default_confidence")]
20    confidence: f32,
21    #[serde(default)]
22    edges: Vec<EdgeParam>,
23}
24
25#[derive(Debug, Deserialize)]
26struct EdgeParam {
27    target_id: u64,
28    edge_type: String,
29    #[serde(default = "default_weight")]
30    weight: f32,
31}
32
33fn default_confidence() -> f32 {
34    0.9
35}
36
37fn default_weight() -> f32 {
38    1.0
39}
40
41/// Return the tool definition for memory_add.
42pub fn definition() -> ToolDefinition {
43    ToolDefinition {
44        name: "memory_add".to_string(),
45        description: Some("Add a new cognitive event to the memory graph".to_string()),
46        input_schema: json!({
47            "type": "object",
48            "properties": {
49                "event_type": {
50                    "type": "string",
51                    "enum": ["fact", "decision", "inference", "correction", "skill", "episode"],
52                    "description": "Type of cognitive event"
53                },
54                "content": {
55                    "type": "string",
56                    "description": "The content of the memory"
57                },
58                "confidence": {
59                    "type": "number",
60                    "minimum": 0.0,
61                    "maximum": 1.0,
62                    "default": 0.9,
63                    "description": "Confidence level (0.0 to 1.0)"
64                },
65                "edges": {
66                    "type": "array",
67                    "items": {
68                        "type": "object",
69                        "properties": {
70                            "target_id": { "type": "integer" },
71                            "edge_type": {
72                                "type": "string",
73                                "enum": ["caused_by", "supports", "contradicts", "supersedes", "related_to", "part_of", "temporal_next"]
74                            },
75                            "weight": { "type": "number", "default": 1.0 }
76                        },
77                        "required": ["target_id", "edge_type"]
78                    }
79                }
80            },
81            "required": ["event_type", "content"]
82        }),
83    }
84}
85
86/// Execute the memory_add tool.
87pub async fn execute(
88    args: Value,
89    session: &Arc<Mutex<SessionManager>>,
90) -> McpResult<ToolCallResult> {
91    let params: AddParams =
92        serde_json::from_value(args).map_err(|e| McpError::InvalidParams(e.to_string()))?;
93
94    // Validate confidence is within [0.0, 1.0]
95    if !(0.0..=1.0).contains(&params.confidence) {
96        return Err(McpError::InvalidParams(format!(
97            "confidence must be between 0.0 and 1.0, got {}",
98            params.confidence
99        )));
100    }
101
102    let event_type = EventType::from_name(&params.event_type).ok_or_else(|| {
103        McpError::InvalidParams(format!("Unknown event type: {}", params.event_type))
104    })?;
105
106    let edges: Vec<(u64, EdgeType, f32)> = params
107        .edges
108        .iter()
109        .map(|e| {
110            let edge_type = EdgeType::from_name(&e.edge_type).ok_or_else(|| {
111                McpError::InvalidParams(format!("Unknown edge type: {}", e.edge_type))
112            })?;
113            Ok((e.target_id, edge_type, e.weight))
114        })
115        .collect::<McpResult<Vec<_>>>()?;
116
117    let mut session = session.lock().await;
118    let (node_id, edges_created) =
119        session.add_event(event_type, &params.content, params.confidence, edges)?;
120
121    Ok(ToolCallResult::json(&json!({
122        "node_id": node_id,
123        "event_type": params.event_type,
124        "edges_created": edges_created
125    })))
126}