Skip to main content

agentic_memory_mcp/tools/
memory_temporal.rs

1//! Tool: memory_temporal — Compare knowledge across time periods.
2
3use std::sync::Arc;
4use tokio::sync::Mutex;
5
6use serde::Deserialize;
7use serde_json::{json, Value};
8
9use agentic_memory::{TemporalParams, TimeRange};
10
11use crate::session::SessionManager;
12use crate::types::{McpError, McpResult, ToolCallResult, ToolDefinition};
13
14#[derive(Debug, Deserialize)]
15struct TemporalInputParams {
16    range_a: RangeSpec,
17    range_b: RangeSpec,
18}
19
20#[derive(Debug, Deserialize)]
21#[serde(tag = "type")]
22enum RangeSpec {
23    #[serde(rename = "time_window")]
24    TimeWindow { start: u64, end: u64 },
25    #[serde(rename = "session")]
26    Session { session_id: u32 },
27    #[serde(rename = "sessions")]
28    Sessions { session_ids: Vec<u32> },
29}
30
31impl RangeSpec {
32    fn to_time_range(&self) -> TimeRange {
33        match self {
34            RangeSpec::TimeWindow { start, end } => TimeRange::TimeWindow {
35                start: *start,
36                end: *end,
37            },
38            RangeSpec::Session { session_id } => TimeRange::Session(*session_id),
39            RangeSpec::Sessions { session_ids } => TimeRange::Sessions(session_ids.clone()),
40        }
41    }
42}
43
44/// Return the tool definition for memory_temporal.
45pub fn definition() -> ToolDefinition {
46    ToolDefinition {
47        name: "memory_temporal".to_string(),
48        description: Some("Compare knowledge across two time periods".to_string()),
49        input_schema: json!({
50            "type": "object",
51            "properties": {
52                "range_a": {
53                    "oneOf": [
54                        {
55                            "type": "object",
56                            "properties": {
57                                "type": { "const": "time_window" },
58                                "start": { "type": "integer", "minimum": 0 },
59                                "end": { "type": "integer", "minimum": 0 }
60                            },
61                            "required": ["type", "start", "end"],
62                            "additionalProperties": false
63                        },
64                        {
65                            "type": "object",
66                            "properties": {
67                                "type": { "const": "session" },
68                                "session_id": { "type": "integer", "minimum": 0 }
69                            },
70                            "required": ["type", "session_id"],
71                            "additionalProperties": false
72                        },
73                        {
74                            "type": "object",
75                            "properties": {
76                                "type": { "const": "sessions" },
77                                "session_ids": { "type": "array", "items": { "type": "integer", "minimum": 0 } }
78                            },
79                            "required": ["type", "session_ids"],
80                            "additionalProperties": false
81                        }
82                    ]
83                },
84                "range_b": {
85                    "description": "Same structure as range_a",
86                    "oneOf": [
87                        {
88                            "type": "object",
89                            "properties": {
90                                "type": { "const": "time_window" },
91                                "start": { "type": "integer", "minimum": 0 },
92                                "end": { "type": "integer", "minimum": 0 }
93                            },
94                            "required": ["type", "start", "end"],
95                            "additionalProperties": false
96                        },
97                        {
98                            "type": "object",
99                            "properties": {
100                                "type": { "const": "session" },
101                                "session_id": { "type": "integer", "minimum": 0 }
102                            },
103                            "required": ["type", "session_id"],
104                            "additionalProperties": false
105                        },
106                        {
107                            "type": "object",
108                            "properties": {
109                                "type": { "const": "sessions" },
110                                "session_ids": { "type": "array", "items": { "type": "integer", "minimum": 0 } }
111                            },
112                            "required": ["type", "session_ids"],
113                            "additionalProperties": false
114                        }
115                    ]
116                }
117            },
118            "required": ["range_a", "range_b"]
119        }),
120    }
121}
122
123/// Execute the memory_temporal tool.
124pub async fn execute(
125    args: Value,
126    session: &Arc<Mutex<SessionManager>>,
127) -> McpResult<ToolCallResult> {
128    let params: TemporalInputParams =
129        serde_json::from_value(args).map_err(|e| McpError::InvalidParams(e.to_string()))?;
130
131    let temporal_params = TemporalParams {
132        range_a: params.range_a.to_time_range(),
133        range_b: params.range_b.to_time_range(),
134    };
135
136    let session = session.lock().await;
137
138    let result = session
139        .query_engine()
140        .temporal(session.graph(), temporal_params)
141        .map_err(|e| McpError::AgenticMemory(format!("Temporal comparison failed: {e}")))?;
142
143    Ok(ToolCallResult::json(&json!({
144        "added": result.added,
145        "corrected": result.corrected,
146        "unchanged": result.unchanged,
147        "potentially_stale": result.potentially_stale,
148        "summary": {
149            "added_count": result.added.len(),
150            "corrected_count": result.corrected.len(),
151            "unchanged_count": result.unchanged.len(),
152            "stale_count": result.potentially_stale.len(),
153        }
154    })))
155}