1use std::future::Future;
5use std::pin::Pin;
6
7use crate::error::Result;
8use crate::types::Message;
9
10pub trait Memory: Send + Sync {
12 fn add_message(&self, message: &Message) -> impl Future<Output = Result<()>> + Send;
19
20 fn get_messages(&self) -> impl Future<Output = Result<Vec<Message>>> + Send;
22
23 fn clear(&self) -> impl Future<Output = Result<()>> + Send;
25}
26
27pub trait ErasedMemory: Send + Sync {
29 fn add_message_erased<'a>(
30 &'a self,
31 message: &'a Message,
32 ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>>;
33
34 fn get_messages_erased<'a>(
35 &'a self,
36 ) -> Pin<Box<dyn Future<Output = Result<Vec<Message>>> + Send + 'a>>;
37
38 fn clear_erased<'a>(&'a self) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>>;
39}
40
41impl<T: Memory> ErasedMemory for T {
42 fn add_message_erased<'a>(
43 &'a self,
44 message: &'a Message,
45 ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>> {
46 Box::pin(self.add_message(message))
47 }
48
49 fn get_messages_erased<'a>(
50 &'a self,
51 ) -> Pin<Box<dyn Future<Output = Result<Vec<Message>>> + Send + 'a>> {
52 Box::pin(self.get_messages())
53 }
54
55 fn clear_erased<'a>(&'a self) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>> {
56 Box::pin(self.clear())
57 }
58}
59
60pub type SharedMemory = std::sync::Arc<dyn ErasedMemory>;
62
63#[cfg(test)]
64mod tests {
65 use super::{Memory, SharedMemory};
66 use crate::{Message, Result, Role};
67 use std::sync::{Arc, Mutex};
68
69 struct VecMemory(Mutex<Vec<Message>>);
71
72 impl Memory for VecMemory {
73 async fn add_message(&self, message: &Message) -> Result<()> {
74 self.0.lock().unwrap().push(message.clone());
75 Ok(())
76 }
77
78 async fn get_messages(&self) -> Result<Vec<Message>> {
79 Ok(self.0.lock().unwrap().clone())
80 }
81
82 async fn clear(&self) -> Result<()> {
83 self.0.lock().unwrap().clear();
84 Ok(())
85 }
86 }
87
88 #[tokio::test]
89 async fn memory_is_implementable_from_core_alone() {
90 let mem = VecMemory(Mutex::new(Vec::new()));
91 mem.add_message(&Message::user("hi")).await.unwrap();
92 assert_eq!(mem.get_messages().await.unwrap().len(), 1);
93 assert_eq!(mem.get_messages().await.unwrap()[0].role, Role::User);
94 mem.clear().await.unwrap();
95 assert!(mem.get_messages().await.unwrap().is_empty());
96
97 let shared: SharedMemory = Arc::new(VecMemory(Mutex::new(Vec::new())));
98 shared
99 .add_message_erased(&Message::user("x"))
100 .await
101 .unwrap();
102 assert_eq!(shared.get_messages_erased().await.unwrap().len(), 1);
103 }
104}