agent/
memory.rs

1use crate::message::Message;
2use crate::error::Result;
3use async_trait::async_trait;
4use std::sync::Arc;
5use tokio::sync::RwLock;
6
7#[async_trait]
8pub trait Memory: Send + Sync {
9    async fn add(&mut self, message: Message) -> Result<()>;
10    async fn get_all(&self) -> Result<Vec<Message>>;
11    async fn get_recent(&self, count: usize) -> Result<Vec<Message>>;
12    async fn clear(&mut self) -> Result<()>;
13}
14
15#[derive(Debug, Clone)]
16pub struct MemoryStore {
17    messages: Arc<RwLock<Vec<Message>>>,
18    max_size: Option<usize>,
19}
20
21impl MemoryStore {
22    pub fn new() -> Self {
23        Self {
24            messages: Arc::new(RwLock::new(Vec::new())),
25            max_size: None,
26        }
27    }
28
29    pub fn with_max_size(max_size: usize) -> Self {
30        Self {
31            messages: Arc::new(RwLock::new(Vec::new())),
32            max_size: Some(max_size),
33        }
34    }
35}
36
37impl Default for MemoryStore {
38    fn default() -> Self {
39        Self::new()
40    }
41}
42
43#[async_trait]
44impl Memory for MemoryStore {
45    async fn add(&mut self, message: Message) -> Result<()> {
46        let mut messages = self.messages.write().await;
47        messages.push(message);
48        
49        if let Some(max_size) = self.max_size {
50            if messages.len() > max_size {
51                let excess = messages.len() - max_size;
52                messages.drain(0..excess);
53            }
54        }
55        
56        Ok(())
57    }
58
59    async fn get_all(&self) -> Result<Vec<Message>> {
60        let messages = self.messages.read().await;
61        Ok(messages.clone())
62    }
63
64    async fn get_recent(&self, count: usize) -> Result<Vec<Message>> {
65        let messages = self.messages.read().await;
66        let start = messages.len().saturating_sub(count);
67        Ok(messages[start..].to_vec())
68    }
69
70    async fn clear(&mut self) -> Result<()> {
71        let mut messages = self.messages.write().await;
72        messages.clear();
73        Ok(())
74    }
75}
agent/memory.rs

agent/
memory.rs
