llm_memory_graph/storage/
serialization.rs

1//! Serialization utilities for storage
2
3use crate::error::{Error, Result};
4use crate::types::{Edge, Node};
5
6/// Serialization format options
7#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8pub enum SerializationFormat {
9    /// JSON format (human-readable, slower)
10    Json,
11    /// MessagePack format (binary, faster)
12    MessagePack,
13    /// Bincode format (binary, fastest)
14    Bincode,
15}
16
17/// Handles serialization and deserialization of graph entities
18pub struct Serializer {
19    format: SerializationFormat,
20}
21
22impl Serializer {
23    /// Create a new serializer with the specified format
24    #[must_use]
25    pub const fn new(format: SerializationFormat) -> Self {
26        Self { format }
27    }
28
29    /// Serialize a node to bytes
30    pub fn serialize_node(&self, node: &Node) -> Result<Vec<u8>> {
31        match self.format {
32            SerializationFormat::Json => {
33                serde_json::to_vec(node).map_err(|e| Error::Serialization(e.to_string()))
34            }
35            SerializationFormat::MessagePack => {
36                rmp_serde::to_vec(node).map_err(|e| Error::Serialization(e.to_string()))
37            }
38            SerializationFormat::Bincode => {
39                bincode::serialize(node).map_err(|e| Error::Serialization(e.to_string()))
40            }
41        }
42    }
43
44    /// Deserialize a node from bytes
45    pub fn deserialize_node(&self, bytes: &[u8]) -> Result<Node> {
46        match self.format {
47            SerializationFormat::Json => {
48                serde_json::from_slice(bytes).map_err(|e| Error::Serialization(e.to_string()))
49            }
50            SerializationFormat::MessagePack => {
51                rmp_serde::from_slice(bytes).map_err(|e| Error::Serialization(e.to_string()))
52            }
53            SerializationFormat::Bincode => {
54                bincode::deserialize(bytes).map_err(|e| Error::Serialization(e.to_string()))
55            }
56        }
57    }
58
59    /// Serialize an edge to bytes
60    pub fn serialize_edge(&self, edge: &Edge) -> Result<Vec<u8>> {
61        match self.format {
62            SerializationFormat::Json => {
63                serde_json::to_vec(edge).map_err(|e| Error::Serialization(e.to_string()))
64            }
65            SerializationFormat::MessagePack => {
66                rmp_serde::to_vec(edge).map_err(|e| Error::Serialization(e.to_string()))
67            }
68            SerializationFormat::Bincode => {
69                bincode::serialize(edge).map_err(|e| Error::Serialization(e.to_string()))
70            }
71        }
72    }
73
74    /// Deserialize an edge from bytes
75    pub fn deserialize_edge(&self, bytes: &[u8]) -> Result<Edge> {
76        match self.format {
77            SerializationFormat::Json => {
78                serde_json::from_slice(bytes).map_err(|e| Error::Serialization(e.to_string()))
79            }
80            SerializationFormat::MessagePack => {
81                rmp_serde::from_slice(bytes).map_err(|e| Error::Serialization(e.to_string()))
82            }
83            SerializationFormat::Bincode => {
84                bincode::deserialize(bytes).map_err(|e| Error::Serialization(e.to_string()))
85            }
86        }
87    }
88}
89
90impl Default for Serializer {
91    fn default() -> Self {
92        Self::new(SerializationFormat::MessagePack)
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99    use crate::types::{NodeId, PromptNode, SessionId};
100
101    #[test]
102    fn test_node_json_serialization() {
103        let session_id = SessionId::new();
104        let prompt = PromptNode::new(session_id, "Test".to_string());
105        let node = Node::Prompt(prompt);
106
107        let serializer = Serializer::new(SerializationFormat::Json);
108        let bytes = serializer.serialize_node(&node).unwrap();
109        let deserialized = serializer.deserialize_node(&bytes).unwrap();
110
111        assert_eq!(node.id(), deserialized.id());
112    }
113
114    #[test]
115    fn test_node_messagepack_serialization() {
116        let session_id = SessionId::new();
117        let prompt = PromptNode::new(session_id, "Test".to_string());
118        let node = Node::Prompt(prompt);
119
120        let serializer = Serializer::new(SerializationFormat::MessagePack);
121        let bytes = serializer.serialize_node(&node).unwrap();
122        let deserialized = serializer.deserialize_node(&bytes).unwrap();
123
124        assert_eq!(node.id(), deserialized.id());
125    }
126
127    #[test]
128    fn test_edge_serialization() {
129        use crate::types::{Edge, EdgeType};
130
131        let from = NodeId::new();
132        let to = NodeId::new();
133        let edge = Edge::new(from, to, EdgeType::Follows);
134
135        let serializer = Serializer::default();
136        let bytes = serializer.serialize_edge(&edge).unwrap();
137        let deserialized = serializer.deserialize_edge(&bytes).unwrap();
138
139        assert_eq!(edge.id, deserialized.id);
140        assert_eq!(edge.edge_type, deserialized.edge_type);
141    }
142}
llm_memory_graph/storage/serialization.rs

llm_memory_graph/storage/
serialization.rs
