Skip to main content

a3s_code_core/context/
fs_provider.rs

1//! File System Context Provider
2//!
3//! Provides simple file-based RAG (Retrieval-Augmented Generation):
4//! - Index files in a directory with glob pattern filtering
5//! - Simple keyword-based search with relevance scoring
6//! - Support for file size limits and exclusion patterns
7
8use crate::context::{ContextItem, ContextProvider, ContextQuery, ContextResult, ContextType};
9use async_trait::async_trait;
10use ignore::WalkBuilder;
11use std::collections::HashMap;
12use std::fs;
13use std::path::{Path, PathBuf};
14use std::sync::Arc;
15use tokio::sync::RwLock;
16
17/// File system context provider configuration
18#[derive(Debug, Clone)]
19pub struct FileSystemContextConfig {
20    /// Root directory to index
21    pub root_path: PathBuf,
22    /// Include patterns (glob syntax: ["**/*.rs", "**/*.md"])
23    pub include_patterns: Vec<String>,
24    /// Exclude patterns (glob syntax: ["**/target/**", "**/node_modules/**"])
25    pub exclude_patterns: Vec<String>,
26    /// Maximum file size in bytes (default: 1MB)
27    pub max_file_size: usize,
28    /// Whether to enable cache (default: true)
29    pub enable_cache: bool,
30}
31
32impl FileSystemContextConfig {
33    /// Create a new config with default settings
34    pub fn new(root_path: impl Into<PathBuf>) -> Self {
35        Self {
36            root_path: root_path.into(),
37            include_patterns: vec!["**/*.rs".to_string(), "**/*.md".to_string()],
38            exclude_patterns: vec![
39                "**/target/**".to_string(),
40                "**/node_modules/**".to_string(),
41                "**/.git/**".to_string(),
42            ],
43            max_file_size: 1024 * 1024, // 1MB
44            enable_cache: true,
45        }
46    }
47
48    /// Set include patterns
49    pub fn with_include_patterns(mut self, patterns: Vec<String>) -> Self {
50        self.include_patterns = patterns;
51        self
52    }
53
54    /// Set exclude patterns
55    pub fn with_exclude_patterns(mut self, patterns: Vec<String>) -> Self {
56        self.exclude_patterns = patterns;
57        self
58    }
59
60    /// Set max file size
61    pub fn with_max_file_size(mut self, size: usize) -> Self {
62        self.max_file_size = size;
63        self
64    }
65
66    /// Enable/disable cache
67    pub fn with_cache(mut self, enable: bool) -> Self {
68        self.enable_cache = enable;
69        self
70    }
71}
72
73/// Indexed file entry
74#[derive(Debug, Clone)]
75struct IndexedFile {
76    path: PathBuf,
77    content: String,
78    size: usize,
79}
80
81/// File system context provider
82pub struct FileSystemContextProvider {
83    config: FileSystemContextConfig,
84    /// Cached indexed files
85    cache: Arc<RwLock<HashMap<PathBuf, IndexedFile>>>,
86}
87
88impl FileSystemContextProvider {
89    /// Create a new file system context provider
90    pub fn new(config: FileSystemContextConfig) -> Self {
91        Self {
92            config,
93            cache: Arc::new(RwLock::new(HashMap::new())),
94        }
95    }
96
97    /// Index files in the root directory
98    async fn index_files(&self) -> anyhow::Result<Vec<IndexedFile>> {
99        let mut files = Vec::new();
100
101        let walker = WalkBuilder::new(&self.config.root_path)
102            .hidden(false)
103            .git_ignore(true)
104            .build();
105
106        for entry in walker {
107            let entry = entry.map_err(|e| anyhow::anyhow!("Walk error: {}", e))?;
108            let path = entry.path();
109
110            if !path.is_file() {
111                continue;
112            }
113
114            let metadata =
115                fs::metadata(path).map_err(|e| anyhow::anyhow!("Metadata error: {}", e))?;
116            if metadata.len() > self.config.max_file_size as u64 {
117                continue;
118            }
119
120            if !self.matches_include_patterns(path) {
121                continue;
122            }
123
124            if self.matches_exclude_patterns(path) {
125                continue;
126            }
127
128            let Some(content) =
129                crate::context::read_utf8_file_bounded(path, self.config.max_file_size)
130                    .map_err(|e| anyhow::anyhow!("Read error: {}", e))?
131            else {
132                continue;
133            };
134
135            files.push(IndexedFile {
136                path: path.to_path_buf(),
137                content,
138                size: metadata.len() as usize,
139            });
140        }
141
142        Ok(files)
143    }
144
145    fn matches_include_patterns(&self, path: &Path) -> bool {
146        if self.config.include_patterns.is_empty() {
147            return true;
148        }
149
150        // Normalize to forward slashes for consistent cross-platform glob matching
151        let path_str = path.to_string_lossy().replace('\\', "/");
152        self.config.include_patterns.iter().any(|pattern| {
153            glob::Pattern::new(pattern)
154                .map(|p| p.matches(&path_str))
155                .unwrap_or(false)
156        })
157    }
158
159    fn matches_exclude_patterns(&self, path: &Path) -> bool {
160        // Normalize to forward slashes for consistent cross-platform glob matching
161        let path_str = path.to_string_lossy().replace('\\', "/");
162        self.config.exclude_patterns.iter().any(|pattern| {
163            glob::Pattern::new(pattern)
164                .map(|p| p.matches(&path_str))
165                .unwrap_or(false)
166        })
167    }
168
169    async fn search_simple(
170        &self,
171        query: &str,
172        files: &[IndexedFile],
173        max_results: usize,
174    ) -> Vec<(IndexedFile, f32)> {
175        let query_lower = query.to_lowercase();
176        let keywords: Vec<&str> = query_lower.split_whitespace().collect();
177
178        let mut results: Vec<(IndexedFile, f32)> = files
179            .iter()
180            .filter_map(|file| {
181                let content_lower = file.content.to_lowercase();
182                let mut score = 0.0;
183                for keyword in &keywords {
184                    let count = content_lower.matches(keyword).count();
185                    score += count as f32;
186                }
187
188                if score > 0.0 {
189                    let normalized_score = score / (file.content.len() as f32).sqrt();
190                    Some((file.clone(), normalized_score))
191                } else {
192                    None
193                }
194            })
195            .collect();
196
197        results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
198        results.truncate(max_results);
199        results
200    }
201
202    async fn update_cache(&self, files: Vec<IndexedFile>) {
203        if !self.config.enable_cache {
204            return;
205        }
206
207        let mut cache = self.cache.write().await;
208        cache.clear();
209        for file in files {
210            cache.insert(file.path.clone(), file);
211        }
212    }
213
214    async fn get_files(&self) -> anyhow::Result<Vec<IndexedFile>> {
215        if self.config.enable_cache {
216            let cache = self.cache.read().await;
217            if !cache.is_empty() {
218                return Ok(cache.values().cloned().collect());
219            }
220        }
221
222        let files = self.index_files().await?;
223        self.update_cache(files.clone()).await;
224        Ok(files)
225    }
226}
227
228#[async_trait]
229impl ContextProvider for FileSystemContextProvider {
230    fn name(&self) -> &str {
231        "filesystem"
232    }
233
234    async fn query(&self, query: &ContextQuery) -> anyhow::Result<ContextResult> {
235        let files = self.get_files().await?;
236        let results = self
237            .search_simple(&query.query, &files, query.max_results)
238            .await;
239
240        let items: Vec<ContextItem> = results
241            .into_iter()
242            .map(|(file, score)| {
243                let content = match query.depth {
244                    crate::context::ContextDepth::Abstract => {
245                        file.content.chars().take(500).collect::<String>()
246                    }
247                    crate::context::ContextDepth::Overview => {
248                        file.content.chars().take(2000).collect::<String>()
249                    }
250                    crate::context::ContextDepth::Full => file.content.clone(),
251                };
252
253                let token_count = content.split_whitespace().count();
254
255                ContextItem::new(
256                    file.path.to_string_lossy().to_string(),
257                    ContextType::Resource,
258                    content,
259                )
260                .with_token_count(token_count)
261                .with_relevance(score)
262                .with_source(format!("file:{}", file.path.display()))
263                .with_provenance("file_system")
264                .with_priority(0.55)
265                .with_trust(0.8)
266                .with_freshness(0.75)
267                .with_metadata("path", serde_json::json!(file.path.to_string_lossy()))
268                .with_metadata("size", serde_json::json!(file.size))
269            })
270            .collect();
271
272        let total_tokens = items
273            .iter()
274            .fold(0usize, |total, item| total.saturating_add(item.token_count));
275        let truncated = items.len() < files.len();
276
277        Ok(ContextResult {
278            items,
279            total_tokens,
280            provider: self.name().to_string(),
281            truncated,
282        })
283    }
284}
285
286#[cfg(test)]
287mod tests {
288    use super::*;
289    use std::fs::File;
290    use std::io::Write;
291    use tempfile::TempDir;
292
293    fn create_test_files(dir: &Path) -> anyhow::Result<()> {
294        let mut file1 = File::create(dir.join("test1.rs"))?;
295        writeln!(file1, "fn main() {{\n    println!(\"Hello, world!\");\n}}")?;
296
297        let mut file2 = File::create(dir.join("test2.md"))?;
298        writeln!(
299            file2,
300            "# Test Document\n\nThis is a test document about Rust programming."
301        )?;
302
303        fs::create_dir(dir.join("subdir"))?;
304        let mut file4 = File::create(dir.join("subdir/test4.rs"))?;
305        writeln!(file4, "fn test() {{\n    // Test function\n}}")?;
306
307        Ok(())
308    }
309
310    #[tokio::test]
311    async fn test_index_files() {
312        let temp_dir = TempDir::new().unwrap();
313        create_test_files(temp_dir.path()).unwrap();
314
315        let config = FileSystemContextConfig::new(temp_dir.path());
316        let provider = FileSystemContextProvider::new(config);
317
318        let files = provider.index_files().await.unwrap();
319        assert!(files.len() >= 2);
320    }
321
322    #[tokio::test]
323    async fn test_search_simple() {
324        let temp_dir = TempDir::new().unwrap();
325        create_test_files(temp_dir.path()).unwrap();
326
327        let config = FileSystemContextConfig::new(temp_dir.path());
328        let provider = FileSystemContextProvider::new(config);
329
330        let query = ContextQuery::new("Rust programming");
331        let result = provider.query(&query).await.unwrap();
332
333        assert!(!result.items.is_empty());
334        assert!(result
335            .items
336            .iter()
337            .any(|item| item.content.contains("Rust")));
338    }
339}