agent_office/storage/
memory.rs

1use crate::domain::{Edge, EdgeId, GraphQuery, Node, NodeId};
2use crate::storage::{EdgeDirection, GraphStorage, Result, StorageError, SearchQuery, SearchResults};
3use async_trait::async_trait;
4use std::collections::HashMap;
5use std::sync::Arc;
6use tokio::sync::RwLock;
7
8#[derive(Clone)]
9pub struct InMemoryStorage {
10    nodes: Arc<RwLock<HashMap<NodeId, Node>>>,
11    edges: Arc<RwLock<HashMap<EdgeId, Edge>>>,
12}
13
14impl InMemoryStorage {
15    pub fn new() -> Self {
16        Self {
17            nodes: Arc::new(RwLock::new(HashMap::new())),
18            edges: Arc::new(RwLock::new(HashMap::new())),
19        }
20    }
21
22    fn matches_query(node: &Node, query: &GraphQuery) -> bool {
23        // Check node type filter
24        if let Some(ref types) = query.node_types {
25            if !types.contains(&node.node_type) {
26                return false;
27            }
28        }
29
30        // Check property filters
31        if let Some(ref filters) = query.property_filters {
32            for (key, expected_value) in filters {
33                match node.properties.get(key) {
34                    Some(actual_value) if actual_value == expected_value => continue,
35                    _ => return false,
36                }
37            }
38        }
39
40        true
41    }
42    
43    fn matches_search_query(node: &Node, query: &SearchQuery) -> bool {
44        // Check node types
45        if !query.node_types.is_empty() && !query.node_types.contains(&node.node_type) {
46            return false;
47        }
48        
49        // Check text search
50        if let Some(ref search_text) = query.search_text {
51            let search_lower = search_text.to_lowercase();
52            let node_text = serde_json::to_string(&node.properties).unwrap_or_default().to_lowercase();
53            if !node_text.contains(&search_lower) {
54                return false;
55            }
56        }
57        
58        // Check created time range
59        if let Some(after) = query.created_after {
60            if node.created_at < after {
61                return false;
62            }
63        }
64        if let Some(before) = query.created_before {
65            if node.created_at > before {
66                return false;
67            }
68        }
69        
70        // Check updated time range
71        if let Some(after) = query.updated_after {
72            if node.updated_at < after {
73                return false;
74            }
75        }
76        
77        // Check property filters
78        for (key, value) in &query.property_filters {
79            match node.properties.get(key) {
80                Some(prop_val) => {
81                    let prop_str = serde_json::to_string(prop_val).unwrap_or_default();
82                    let value_str = format!("\"{}\"", value);
83                    if prop_str != value_str && prop_str != *value {
84                        return false;
85                    }
86                }
87                None => return false,
88            }
89        }
90        
91        true
92    }
93}
94
95impl Default for InMemoryStorage {
96    fn default() -> Self {
97        Self::new()
98    }
99}
100
101#[async_trait]
102impl GraphStorage for InMemoryStorage {
103    async fn create_node(&self, node: &Node) -> Result<Node> {
104        let mut nodes = self.nodes.write().await;
105        if nodes.contains_key(&node.id) {
106            return Err(StorageError::ConstraintViolation(
107                format!("Node with ID {} already exists", node.id)
108            ));
109        }
110        nodes.insert(node.id, node.clone());
111        Ok(node.clone())
112    }
113
114    async fn get_node(&self, id: NodeId) -> Result<Node> {
115        let nodes = self.nodes.read().await;
116        nodes.get(&id)
117            .cloned()
118            .ok_or(StorageError::NodeNotFound(id))
119    }
120
121    async fn update_node(&self, node: &Node) -> Result<Node> {
122        let mut nodes = self.nodes.write().await;
123        if !nodes.contains_key(&node.id) {
124            return Err(StorageError::NodeNotFound(node.id));
125        }
126        nodes.insert(node.id, node.clone());
127        Ok(node.clone())
128    }
129
130    async fn delete_node(&self, id: NodeId) -> Result<()> {
131        let mut nodes = self.nodes.write().await;
132        let mut edges = self.edges.write().await;
133        
134        if !nodes.contains_key(&id) {
135            return Err(StorageError::NodeNotFound(id));
136        }
137        
138        // Remove all edges connected to this node
139        edges.retain(|_, edge| {
140            edge.from_node_id != id && edge.to_node_id != id
141        });
142        
143        nodes.remove(&id);
144        Ok(())
145    }
146
147    async fn query_nodes(&self, query: &GraphQuery) -> Result<Vec<Node>> {
148        let nodes = self.nodes.read().await;
149        let mut results: Vec<Node> = nodes
150            .values()
151            .filter(|node| Self::matches_query(node, query))
152            .cloned()
153            .collect();
154        
155        if let Some(limit) = query.limit {
156            results.truncate(limit);
157        }
158        
159        Ok(results)
160    }
161
162    async fn create_edge(&self, edge: &Edge) -> Result<Edge> {
163        let nodes = self.nodes.read().await;
164        
165        // Verify both nodes exist
166        if !nodes.contains_key(&edge.from_node_id) {
167            return Err(StorageError::NodeNotFound(edge.from_node_id));
168        }
169        if !nodes.contains_key(&edge.to_node_id) {
170            return Err(StorageError::NodeNotFound(edge.to_node_id));
171        }
172        
173        drop(nodes);
174        
175        let mut edges = self.edges.write().await;
176        edges.insert(edge.id, edge.clone());
177        Ok(edge.clone())
178    }
179
180    async fn get_edges_from(&self, node_id: NodeId, edge_type: Option<&str>) -> Result<Vec<Edge>> {
181        let edges = self.edges.read().await;
182        let results: Vec<Edge> = edges
183            .values()
184            .filter(|edge| {
185                edge.from_node_id == node_id &&
186                edge_type.map_or(true, |et| edge.edge_type == et)
187            })
188            .cloned()
189            .collect();
190        Ok(results)
191    }
192
193    async fn get_edges_to(&self, node_id: NodeId, edge_type: Option<&str>) -> Result<Vec<Edge>> {
194        let edges = self.edges.read().await;
195        let results: Vec<Edge> = edges
196            .values()
197            .filter(|edge| {
198                edge.to_node_id == node_id &&
199                edge_type.map_or(true, |et| edge.edge_type == et)
200            })
201            .cloned()
202            .collect();
203        Ok(results)
204    }
205
206    async fn get_neighbors(
207        &self,
208        node_id: NodeId,
209        edge_type: Option<&str>,
210        direction: EdgeDirection,
211    ) -> Result<Vec<Node>> {
212        let edges = self.edges.read().await;
213        let nodes = self.nodes.read().await;
214        
215        let mut neighbor_ids: Vec<NodeId> = Vec::new();
216        
217        for edge in edges.values() {
218            let matches_type = edge_type.map_or(true, |et| edge.edge_type == et);
219            
220            match direction {
221                EdgeDirection::Outgoing if edge.from_node_id == node_id && matches_type => {
222                    neighbor_ids.push(edge.to_node_id);
223                }
224                EdgeDirection::Incoming if edge.to_node_id == node_id && matches_type => {
225                    neighbor_ids.push(edge.from_node_id);
226                }
227                _ => {}
228            }
229        }
230        
231        let neighbors: Vec<Node> = neighbor_ids
232            .into_iter()
233            .filter_map(|id| nodes.get(&id).cloned())
234            .collect();
235        
236        Ok(neighbors)
237    }
238    
239    async fn search_nodes(&self, query: &SearchQuery) -> Result<SearchResults<Node>> {
240        let nodes = self.nodes.read().await;
241        
242        // Filter nodes based on query criteria
243        let results: Vec<Node> = nodes.values()
244            .filter(|node| Self::matches_search_query(node, query))
245            .cloned()
246            .collect();
247        
248        // Apply pagination
249        let offset = query.offset;
250        let limit = query.limit;
251        
252        let paginated: Vec<Node> = results.into_iter()
253            .skip(offset)
254            .take(limit)
255            .collect();
256        
257        Ok(SearchResults {
258            items: paginated,
259        })
260    }
261}
262
263#[cfg(test)]
264mod tests {
265    use super::*;
266    use crate::domain::Properties;
267
268    #[tokio::test]
269    async fn test_create_and_get_node() {
270        let storage = InMemoryStorage::new();
271        let node = Node::new("test", Properties::new());
272        
273        let created = storage.create_node(&node).await.unwrap();
274        assert_eq!(created.id, node.id);
275        
276        let retrieved = storage.get_node(node.id).await.unwrap();
277        assert_eq!(retrieved.id, node.id);
278    }
279
280    #[tokio::test]
281    async fn test_create_edge_between_nodes() {
282        let storage = InMemoryStorage::new();
283        
284        let node1 = Node::new("agent", Properties::new());
285        let node2 = Node::new("mailbox", Properties::new());
286        
287        storage.create_node(&node1).await.unwrap();
288        storage.create_node(&node2).await.unwrap();
289        
290        let edge = Edge::new("owns", node1.id, node2.id, Properties::new());
291        let created = storage.create_edge(&edge).await.unwrap();
292        
293        assert_eq!(created.from_node_id, node1.id);
294        assert_eq!(created.to_node_id, node2.id);
295    }
296
297    #[tokio::test]
298    async fn test_query_nodes_with_type_filter() {
299        let storage = InMemoryStorage::new();
300        
301        let agent = Node::new("agent", Properties::new());
302        let mailbox = Node::new("mailbox", Properties::new());
303        
304        storage.create_node(&agent).await.unwrap();
305        storage.create_node(&mailbox).await.unwrap();
306        
307        let query = GraphQuery::new().with_node_type("agent");
308        let results = storage.query_nodes(&query).await.unwrap();
309        
310        assert_eq!(results.len(), 1);
311        assert_eq!(results[0].node_type, "agent");
312    }
313
314    #[tokio::test]
315    async fn test_get_neighbors() {
316        let storage = InMemoryStorage::new();
317        
318        let agent = Node::new("agent", Properties::new());
319        let mailbox1 = Node::new("mailbox", Properties::new());
320        let mailbox2 = Node::new("mailbox", Properties::new());
321        
322        storage.create_node(&agent).await.unwrap();
323        storage.create_node(&mailbox1).await.unwrap();
324        storage.create_node(&mailbox2).await.unwrap();
325        
326        let edge1 = Edge::new("owns", agent.id, mailbox1.id, Properties::new());
327        let edge2 = Edge::new("owns", agent.id, mailbox2.id, Properties::new());
328        
329        storage.create_edge(&edge1).await.unwrap();
330        storage.create_edge(&edge2).await.unwrap();
331        
332        let neighbors = storage.get_neighbors(agent.id, Some("owns"), EdgeDirection::Outgoing).await.unwrap();
333        assert_eq!(neighbors.len(), 2);
334    }
335}
agent_office/storage/memory.rs

agent_office/storage/
memory.rs
