agentic_vision/
storage.rs

1//! .avis binary file format reader/writer for visual memory.
2
3use std::io::{Read, Write};
4use std::path::Path;
5
6use crate::types::{VisionError, VisionResult, VisualMemoryStore, VisualObservation};
7
8/// Magic bytes: "AVIS"
9const AVIS_MAGIC: u32 = 0x41564953;
10
11/// Current format version.
12const FORMAT_VERSION: u16 = 1;
13
14/// Header size in bytes.
15const HEADER_SIZE: usize = 64;
16
17/// Writer for .avis files.
18pub struct AvisWriter;
19
20/// Reader for .avis files.
21pub struct AvisReader;
22
23impl AvisWriter {
24    /// Write a visual memory store to a file.
25    pub fn write_to_file(store: &VisualMemoryStore, path: &Path) -> VisionResult<()> {
26        if let Some(parent) = path.parent() {
27            std::fs::create_dir_all(parent)?;
28        }
29
30        let mut file = std::fs::File::create(path)?;
31        Self::write_to(store, &mut file)
32    }
33
34    /// Write a visual memory store to any writer.
35    pub fn write_to<W: Write>(store: &VisualMemoryStore, writer: &mut W) -> VisionResult<()> {
36        // Serialize observations as JSON (simple, correct, can optimize later)
37        let payload = serde_json::to_vec(&SerializedStore {
38            observations: &store.observations,
39            embedding_dim: store.embedding_dim,
40            next_id: store.next_id,
41            session_count: store.session_count,
42            created_at: store.created_at,
43            updated_at: store.updated_at,
44        })
45        .map_err(|e| VisionError::Storage(format!("Serialization failed: {e}")))?;
46
47        // Write header
48        let mut header = [0u8; HEADER_SIZE];
49        write_u32(&mut header[0..4], AVIS_MAGIC);
50        write_u16(&mut header[4..6], FORMAT_VERSION);
51        write_u16(&mut header[6..8], 0); // flags
52        write_u64(&mut header[8..16], store.observations.len() as u64);
53        write_u32(&mut header[16..20], store.embedding_dim);
54        write_u32(&mut header[20..24], store.session_count);
55        write_u64(&mut header[24..32], store.created_at);
56        write_u64(&mut header[32..40], store.updated_at);
57        write_u64(&mut header[40..48], payload.len() as u64); // payload length
58
59        writer.write_all(&header)?;
60        writer.write_all(&payload)?;
61
62        Ok(())
63    }
64}
65
66impl AvisReader {
67    /// Read a visual memory store from a file.
68    pub fn read_from_file(path: &Path) -> VisionResult<VisualMemoryStore> {
69        let mut file = std::fs::File::open(path)?;
70        Self::read_from(&mut file)
71    }
72
73    /// Read a visual memory store from any reader.
74    pub fn read_from<R: Read>(reader: &mut R) -> VisionResult<VisualMemoryStore> {
75        // Read header
76        let mut header = [0u8; HEADER_SIZE];
77        reader.read_exact(&mut header)?;
78
79        let magic = read_u32(&header[0..4]);
80        if magic != AVIS_MAGIC {
81            return Err(VisionError::Storage(format!(
82                "Invalid magic: expected 0x{AVIS_MAGIC:08X}, got 0x{magic:08X}"
83            )));
84        }
85
86        let version = read_u16(&header[4..6]);
87        if version != FORMAT_VERSION {
88            return Err(VisionError::Storage(format!(
89                "Unsupported version: {version}"
90            )));
91        }
92
93        let _observation_count = read_u64(&header[8..16]);
94        let embedding_dim = read_u32(&header[16..20]);
95        let session_count = read_u32(&header[20..24]);
96        let created_at = read_u64(&header[24..32]);
97        let updated_at = read_u64(&header[32..40]);
98        let payload_len = read_u64(&header[40..48]) as usize;
99
100        // Read payload
101        let mut payload = vec![0u8; payload_len];
102        reader.read_exact(&mut payload)?;
103
104        let serialized: DeserializedStore = serde_json::from_slice(&payload)
105            .map_err(|e| VisionError::Storage(format!("Deserialization failed: {e}")))?;
106
107        let next_id = serialized.next_id;
108
109        Ok(VisualMemoryStore {
110            observations: serialized.observations,
111            embedding_dim,
112            next_id,
113            session_count,
114            created_at,
115            updated_at,
116        })
117    }
118}
119
120#[derive(serde::Serialize)]
121struct SerializedStore<'a> {
122    observations: &'a [VisualObservation],
123    embedding_dim: u32,
124    next_id: u64,
125    session_count: u32,
126    created_at: u64,
127    updated_at: u64,
128}
129
130#[derive(serde::Deserialize)]
131struct DeserializedStore {
132    observations: Vec<VisualObservation>,
133    #[allow(dead_code)]
134    embedding_dim: u32,
135    next_id: u64,
136    #[allow(dead_code)]
137    session_count: u32,
138    #[allow(dead_code)]
139    created_at: u64,
140    #[allow(dead_code)]
141    updated_at: u64,
142}
143
144// Little-endian byte helpers
145fn write_u16(buf: &mut [u8], val: u16) {
146    buf[..2].copy_from_slice(&val.to_le_bytes());
147}
148fn write_u32(buf: &mut [u8], val: u32) {
149    buf[..4].copy_from_slice(&val.to_le_bytes());
150}
151fn write_u64(buf: &mut [u8], val: u64) {
152    buf[..8].copy_from_slice(&val.to_le_bytes());
153}
154fn read_u16(buf: &[u8]) -> u16 {
155    u16::from_le_bytes([buf[0], buf[1]])
156}
157fn read_u32(buf: &[u8]) -> u32 {
158    u32::from_le_bytes([buf[0], buf[1], buf[2], buf[3]])
159}
160fn read_u64(buf: &[u8]) -> u64 {
161    u64::from_le_bytes([buf[0], buf[1], buf[2], buf[3], buf[4], buf[5], buf[6], buf[7]])
162}
163
164#[cfg(test)]
165mod tests {
166    use super::*;
167    use crate::types::{CaptureSource, ObservationMeta};
168
169    fn make_test_observation(id: u64) -> VisualObservation {
170        VisualObservation {
171            id,
172            timestamp: 1708345678,
173            session_id: 1,
174            source: CaptureSource::File {
175                path: "/test/image.png".to_string(),
176            },
177            embedding: vec![0.1, 0.2, 0.3],
178            thumbnail: vec![0xFF, 0xD8, 0xFF],
179            metadata: ObservationMeta {
180                width: 512,
181                height: 512,
182                original_width: 1920,
183                original_height: 1080,
184                labels: vec!["test".to_string()],
185                description: Some("Test observation".to_string()),
186            },
187            memory_link: None,
188        }
189    }
190
191    #[test]
192    fn test_roundtrip_empty() {
193        let store = VisualMemoryStore::new(512);
194        let mut buf = Vec::new();
195        AvisWriter::write_to(&store, &mut buf).unwrap();
196
197        let loaded = AvisReader::read_from(&mut &buf[..]).unwrap();
198        assert_eq!(loaded.count(), 0);
199        assert_eq!(loaded.embedding_dim, 512);
200    }
201
202    #[test]
203    fn test_roundtrip_with_observations() {
204        let mut store = VisualMemoryStore::new(512);
205        store.add(make_test_observation(0));
206        store.add(make_test_observation(0));
207
208        let mut buf = Vec::new();
209        AvisWriter::write_to(&store, &mut buf).unwrap();
210
211        let loaded = AvisReader::read_from(&mut &buf[..]).unwrap();
212        assert_eq!(loaded.count(), 2);
213        assert_eq!(loaded.observations[0].id, 1);
214        assert_eq!(loaded.observations[1].id, 2);
215    }
216
217    #[test]
218    fn test_invalid_magic() {
219        let mut buf = [0u8; HEADER_SIZE + 10];
220        buf[0..4].copy_from_slice(&[0x00, 0x00, 0x00, 0x00]);
221        let result = AvisReader::read_from(&mut &buf[..]);
222        assert!(result.is_err());
223    }
224
225    #[test]
226    fn test_file_roundtrip() {
227        let dir = tempfile::tempdir().unwrap();
228        let path = dir.path().join("test.avis");
229
230        let mut store = VisualMemoryStore::new(512);
231        store.add(make_test_observation(0));
232
233        AvisWriter::write_to_file(&store, &path).unwrap();
234        let loaded = AvisReader::read_from_file(&path).unwrap();
235        assert_eq!(loaded.count(), 1);
236    }
237}
agentic_vision/storage.rs

agentic_vision/
storage.rs
