ai_agents_core/traits/
memory.rs

1//! Memory trait for conversation storage
2
3use async_trait::async_trait;
4
5use crate::error::Result;
6use crate::message::ChatMessage;
7
8/// Snapshot of memory state for persistence
9#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, Default)]
10pub struct MemorySnapshot {
11    #[serde(default)]
12    pub messages: Vec<ChatMessage>,
13    #[serde(default)]
14    pub summary: Option<String>,
15    #[serde(default)]
16    pub summarized_count: usize,
17}
18
19impl MemorySnapshot {
20    pub fn new(messages: Vec<ChatMessage>) -> Self {
21        Self {
22            messages,
23            summary: None,
24            summarized_count: 0,
25        }
26    }
27
28    pub fn with_summary(mut self, summary: String) -> Self {
29        self.summary = Some(summary);
30        self
31    }
32
33    pub fn with_summarized_count(mut self, summarized_count: usize) -> Self {
34        self.summarized_count = summarized_count;
35        self
36    }
37}
38
39/// Core memory trait for storing conversation history.
40///
41/// Built-in implementations: `InMemoryStore` (simple) and `CompactingMemory`
42/// (with LLM-based summarization). Implement this for custom storage strategies.
43#[async_trait]
44pub trait Memory: Send + Sync {
45    /// Append a message to conversation history.
46    async fn add_message(&self, message: ChatMessage) -> Result<()>;
47    /// Get messages. `Some(n)` returns the most recent N messages.
48    async fn get_messages(&self, limit: Option<usize>) -> Result<Vec<ChatMessage>>;
49    /// Remove all messages and reset state.
50    async fn clear(&self) -> Result<()>;
51    /// Number of messages currently stored.
52    fn len(&self) -> usize;
53
54    /// Returns `true` if no messages are stored.
55    fn is_empty(&self) -> bool {
56        self.len() == 0
57    }
58
59    /// Serialize current state for persistence.
60    async fn snapshot(&self) -> Result<MemorySnapshot> {
61        Ok(MemorySnapshot::new(self.get_messages(None).await?))
62    }
63
64    /// Restore from a previously saved snapshot, replacing current state.
65    async fn restore(&self, snapshot: MemorySnapshot) -> Result<()>;
66
67    /// Remove the oldest N messages. Returns empty vec by default.
68    async fn evict_oldest(&self, _count: usize) -> Result<Vec<ChatMessage>> {
69        Ok(vec![])
70    }
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76
77    #[test]
78    fn old_snapshot_defaults_summarized_count() {
79        let json = r#"{"messages":[],"summary":"existing"}"#;
80        let snapshot: MemorySnapshot = serde_json::from_str(json).unwrap();
81        assert_eq!(snapshot.summary.as_deref(), Some("existing"));
82        assert_eq!(snapshot.summarized_count, 0);
83    }
84
85    #[test]
86    fn snapshot_roundtrip_preserves_summarized_count() {
87        let snapshot = MemorySnapshot::new(Vec::new())
88            .with_summary("existing".to_string())
89            .with_summarized_count(12);
90        let json = serde_json::to_string(&snapshot).unwrap();
91        let restored: MemorySnapshot = serde_json::from_str(&json).unwrap();
92        assert_eq!(restored.summary.as_deref(), Some("existing"));
93        assert_eq!(restored.summarized_count, 12);
94    }
95}
ai_agents_core/traits/memory.rs

ai_agents_core/traits/
memory.rs
