agentic_memory/format/
reader.rs

1//! Reads .amem files into in-memory graph.
2
3use std::io::Read;
4use std::path::Path;
5
6use crate::graph::MemoryGraph;
7use crate::types::error::{AmemError, AmemResult};
8use crate::types::header::FileHeader;
9use crate::types::{CognitiveEvent, Edge, EdgeType, EventType};
10
11use super::compression::decompress_content;
12
13/// Reader for .amem binary files.
14pub struct AmemReader;
15
16impl AmemReader {
17    /// Read an .amem file into a MemoryGraph.
18    pub fn read_from_file(path: &Path) -> AmemResult<MemoryGraph> {
19        let data = std::fs::read(path)?;
20        let mut cursor = std::io::Cursor::new(data);
21        Self::read_from(&mut cursor)
22    }
23
24    /// Read from any reader into a MemoryGraph.
25    pub fn read_from(reader: &mut impl Read) -> AmemResult<MemoryGraph> {
26        // Read all data into a buffer
27        let mut data = Vec::new();
28        reader.read_to_end(&mut data)?;
29
30        if data.len() < 64 {
31            return Err(AmemError::Truncated);
32        }
33
34        // Parse header
35        let header = FileHeader::read_from(&mut std::io::Cursor::new(&data[..64]))?;
36
37        let dimension = header.dimension as usize;
38        let node_count = header.node_count as usize;
39        let edge_count = header.edge_count as usize;
40
41        // Read node table
42        let node_table_start = header.node_table_offset as usize;
43        let mut nodes: Vec<CognitiveEvent> = Vec::with_capacity(node_count);
44        let mut node_content_info: Vec<(u64, u32)> = Vec::with_capacity(node_count);
45
46        for i in 0..node_count {
47            let offset = node_table_start + i * 72;
48            if offset + 72 > data.len() {
49                return Err(AmemError::Truncated);
50            }
51            let record = &data[offset..offset + 72];
52            let (event, content_offset, content_length) = parse_node_record(record)?;
53            node_content_info.push((content_offset, content_length));
54            nodes.push(event);
55        }
56
57        // Read edge table
58        let edge_table_start = header.edge_table_offset as usize;
59        let mut edges: Vec<Edge> = Vec::with_capacity(edge_count);
60
61        for i in 0..edge_count {
62            let offset = edge_table_start + i * 32;
63            if offset + 32 > data.len() {
64                return Err(AmemError::Truncated);
65            }
66            let record = &data[offset..offset + 32];
67            edges.push(parse_edge_record(record)?);
68        }
69
70        // Read content block
71        let content_block_start = header.content_block_offset as usize;
72        for (i, node) in nodes.iter_mut().enumerate() {
73            let (content_offset, content_length) = node_content_info[i];
74            if content_length > 0 {
75                let start = content_block_start + content_offset as usize;
76                let end = start + content_length as usize;
77                if end > data.len() {
78                    return Err(AmemError::Truncated);
79                }
80                node.content = decompress_content(&data[start..end])?;
81            }
82        }
83
84        // Read feature vectors
85        let fv_start = header.feature_vec_offset as usize;
86        for (i, node) in nodes.iter_mut().enumerate() {
87            let offset = fv_start + i * dimension * 4;
88            if offset + dimension * 4 > data.len() {
89                return Err(AmemError::Truncated);
90            }
91            let mut vec = Vec::with_capacity(dimension);
92            for j in 0..dimension {
93                let byte_offset = offset + j * 4;
94                let bytes: [u8; 4] = data[byte_offset..byte_offset + 4].try_into().unwrap();
95                vec.push(f32::from_le_bytes(bytes));
96            }
97            node.feature_vec = vec;
98        }
99
100        // Build graph from parts (this rebuilds indexes)
101        MemoryGraph::from_parts(nodes, edges, dimension)
102    }
103}
104
105/// Parse a 72-byte node record.
106fn parse_node_record(data: &[u8]) -> AmemResult<(CognitiveEvent, u64, u32)> {
107    let id = u64::from_le_bytes(data[0..8].try_into().unwrap());
108    let event_type_byte = data[8];
109    let event_type = EventType::from_u8(event_type_byte).ok_or(AmemError::Corrupt(0))?;
110    // bytes 9..12: padding
111    let created_at = u64::from_le_bytes(data[12..20].try_into().unwrap());
112    let session_id = u32::from_le_bytes(data[20..24].try_into().unwrap());
113    let confidence = f32::from_le_bytes(data[24..28].try_into().unwrap());
114    let access_count = u32::from_le_bytes(data[28..32].try_into().unwrap());
115    let last_accessed = u64::from_le_bytes(data[32..40].try_into().unwrap());
116    let decay_score = f32::from_le_bytes(data[40..44].try_into().unwrap());
117    let content_offset = u64::from_le_bytes(data[44..52].try_into().unwrap());
118    let content_length = u32::from_le_bytes(data[52..56].try_into().unwrap());
119    // edge_offset at 56..64 (not needed for in-memory construction)
120    // edge_count at 64..66 (not needed)
121    // padding at 66..72 (not needed)
122
123    let event = CognitiveEvent {
124        id,
125        event_type,
126        created_at,
127        session_id,
128        confidence,
129        access_count,
130        last_accessed,
131        decay_score,
132        content: String::new(),  // Will be filled from content block
133        feature_vec: Vec::new(), // Will be filled from feature vec block
134    };
135
136    Ok((event, content_offset, content_length))
137}
138
139/// Parse a 32-byte edge record.
140fn parse_edge_record(data: &[u8]) -> AmemResult<Edge> {
141    let source_id = u64::from_le_bytes(data[0..8].try_into().unwrap());
142    let target_id = u64::from_le_bytes(data[8..16].try_into().unwrap());
143    let edge_type_byte = data[16];
144    let edge_type = EdgeType::from_u8(edge_type_byte).ok_or(AmemError::Corrupt(0))?;
145    // bytes 17..20: padding
146    let weight = f32::from_le_bytes(data[20..24].try_into().unwrap());
147    let created_at = u64::from_le_bytes(data[24..32].try_into().unwrap());
148
149    Ok(Edge {
150        source_id,
151        target_id,
152        edge_type,
153        weight,
154        created_at,
155    })
156}
agentic_memory/format/reader.rs

agentic_memory/format/
reader.rs
