semtree_rag/
indexer.rs

1use std::path::{Path, PathBuf};
2use std::sync::Arc;
3
4use semtree_embed::Embedder;
5use semtree_parse::extract_file;
6use semtree_store::VectorStore;
7use tracing::{debug, warn};
8
9use crate::{ChunkRegistry, FileManifest, RagError};
10
11pub struct Indexer {
12    embedder: Arc<dyn Embedder>,
13    store: Arc<dyn VectorStore>,
14}
15
16impl Indexer {
17    pub fn new(embedder: Arc<dyn Embedder>, store: Arc<dyn VectorStore>) -> Self {
18        Self { embedder, store }
19    }
20
21    pub async fn index_file(
22        &self,
23        path: &Path,
24        registry: &mut ChunkRegistry,
25        manifest: Option<&mut FileManifest>,
26    ) -> Result<usize, RagError> {
27        let content = match std::fs::read_to_string(path) {
28            Ok(c) => c,
29            Err(e) => {
30                warn!("skipping {}: {e}", path.display());
31                return Ok(0);
32            }
33        };
34
35        if let Some(ref manifest) = manifest {
36            if !manifest.is_changed(path, &content) {
37                debug!("unchanged, skipping {}", path.display());
38                return Ok(0);
39            }
40        }
41
42        // Remove stale chunks for this file before re-indexing
43        if let Some(ref manifest) = manifest {
44            for old_id in manifest.chunk_ids(path) {
45                registry.remove(old_id);
46                let _ = self.store.delete(old_id).await;
47            }
48        }
49
50        let mut chunks = match extract_file(path) {
51            Ok(c) => c,
52            Err(e) => {
53                warn!("skipping {}: {e}", path.display());
54                return Ok(0);
55            }
56        };
57
58        if chunks.is_empty() {
59            if let Some(manifest) = manifest {
60                manifest.record(path.to_path_buf(), &content, vec![]);
61            }
62            return Ok(0);
63        }
64
65        for chunk in &mut chunks {
66            chunk.path = path.to_path_buf();
67        }
68
69        let texts: Vec<&str> = chunks.iter().map(|c| c.content.as_str()).collect();
70        let embeddings = self.embedder.embed(&texts).await?;
71
72        let mut chunk_ids = Vec::with_capacity(chunks.len());
73        for (chunk, embedding) in chunks.iter().zip(embeddings.iter()) {
74            self.store.insert(&chunk.id, embedding).await?;
75            registry.insert(chunk.clone());
76            chunk_ids.push(chunk.id.clone());
77            debug!("indexed {}", chunk.name.as_deref().unwrap_or(&chunk.id));
78        }
79
80        if let Some(manifest) = manifest {
81            manifest.record(path.to_path_buf(), &content, chunk_ids);
82        }
83
84        Ok(chunks.len())
85    }
86
87    /// Index a directory, skipping unchanged files when a manifest is provided.
88    /// Calls `on_progress(files_done, total_files)` after each file.
89    pub async fn index_dir(
90        &self,
91        dir: &Path,
92        registry: &mut ChunkRegistry,
93        manifest: Option<&mut FileManifest>,
94        on_progress: impl Fn(usize, usize),
95    ) -> Result<usize, RagError> {
96        let paths = collect_indexable_files(dir);
97        let total = paths.len();
98        let mut total_chunks = 0;
99
100        // Collect manifest as Option<&mut FileManifest> once, then use it per-file.
101        // We can't borrow it mutably in a loop via the outer Option easily, so we
102        // take ownership and pass per-file slices of info manually.
103        let mut manifest = manifest;
104
105        for (done, path) in paths.iter().enumerate() {
106            total_chunks += self
107                .index_file(path, registry, manifest.as_deref_mut())
108                .await?;
109            on_progress(done + 1, total);
110        }
111
112        // Remove manifest entries for files that no longer exist
113        if let Some(ref mut m) = manifest {
114            let stale: Vec<PathBuf> = m.paths().filter(|p| !p.exists()).cloned().collect();
115            for path in stale {
116                for old_id in m.chunk_ids(&path).to_vec() {
117                    registry.remove(&old_id);
118                    let _ = self.store.delete(&old_id).await;
119                }
120                m.remove(&path);
121            }
122        }
123
124        Ok(total_chunks)
125    }
126}
127
128/// Collect all indexable files under `dir`, respecting the ignore list.
129pub fn collect_indexable_files(dir: &Path) -> Vec<PathBuf> {
130    let mut paths = Vec::new();
131    walkdir(dir, &mut paths);
132    paths
133}
134
135const IGNORED_DIRS: &[&str] = &[
136    // dependencies
137    "node_modules",
138    "vendor",
139    ".pnp",
140    // build output
141    "target",
142    "dist",
143    "build",
144    "out",
145    ".next",
146    ".nuxt",
147    "__pycache__",
148    ".pytest_cache",
149    ".mypy_cache",
150    ".ruff_cache",
151    // venvs
152    ".venv",
153    "venv",
154    "env",
155    ".env",
156    // version control
157    ".git",
158    ".svn",
159    ".hg",
160    // index itself
161    ".semtree",
162    // misc
163    ".idea",
164    ".vscode",
165    "coverage",
166    ".turbo",
167    ".cache",
168];
169
170fn is_ignored(dir_name: &str) -> bool {
171    IGNORED_DIRS.contains(&dir_name)
172}
173
174fn walkdir(dir: &Path, paths: &mut Vec<PathBuf>) {
175    if let Ok(entries) = std::fs::read_dir(dir) {
176        for entry in entries.flatten() {
177            let path = entry.path();
178            if path.is_dir() {
179                let name = path.file_name().and_then(|n| n.to_str()).unwrap_or("");
180                if !is_ignored(name) {
181                    walkdir(&path, paths);
182                }
183            } else {
184                paths.push(path);
185            }
186        }
187    }
188}
semtree_rag/indexer.rs

semtree_rag/
indexer.rs
