Skip to main content

semtree_rag/
registry.rs

1use std::collections::HashMap;
2use std::path::Path;
3
4use semtree_core::Chunk;
5
6use crate::RagError;
7
8/// Persists chunk metadata alongside the vector index.
9#[derive(Default)]
10pub struct ChunkRegistry {
11    chunks: HashMap<String, Chunk>,
12}
13
14impl ChunkRegistry {
15    pub fn insert(&mut self, chunk: Chunk) {
16        self.chunks.insert(chunk.id.clone(), chunk);
17    }
18
19    pub fn get(&self, id: &str) -> Option<&Chunk> {
20        self.chunks.get(id)
21    }
22
23    pub fn remove(&mut self, id: &str) -> Option<Chunk> {
24        self.chunks.remove(id)
25    }
26
27    pub fn iter(&self) -> impl Iterator<Item = &Chunk> {
28        self.chunks.values()
29    }
30
31    pub fn len(&self) -> usize {
32        self.chunks.len()
33    }
34
35    pub fn is_empty(&self) -> bool {
36        self.chunks.is_empty()
37    }
38
39    pub fn save(&self, path: &Path) -> Result<(), RagError> {
40        let data = serde_json::to_string(&self.chunks)
41            .map_err(|e| RagError::Io(std::io::Error::other(e)))?;
42        std::fs::write(path.join("chunks.json"), data)?;
43        Ok(())
44    }
45
46    pub fn load(&mut self, path: &Path) -> Result<(), RagError> {
47        let raw = std::fs::read_to_string(path.join("chunks.json"))?;
48        self.chunks =
49            serde_json::from_str(&raw).map_err(|e| RagError::Io(std::io::Error::other(e)))?;
50        Ok(())
51    }
52}