a3s_code_core/context/
fs_provider.rs1use 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#[derive(Debug, Clone)]
19pub struct FileSystemContextConfig {
20 pub root_path: PathBuf,
22 pub include_patterns: Vec<String>,
24 pub exclude_patterns: Vec<String>,
26 pub max_file_size: usize,
28 pub enable_cache: bool,
30}
31
32impl FileSystemContextConfig {
33 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, enable_cache: true,
45 }
46 }
47
48 pub fn with_include_patterns(mut self, patterns: Vec<String>) -> Self {
50 self.include_patterns = patterns;
51 self
52 }
53
54 pub fn with_exclude_patterns(mut self, patterns: Vec<String>) -> Self {
56 self.exclude_patterns = patterns;
57 self
58 }
59
60 pub fn with_max_file_size(mut self, size: usize) -> Self {
62 self.max_file_size = size;
63 self
64 }
65
66 pub fn with_cache(mut self, enable: bool) -> Self {
68 self.enable_cache = enable;
69 self
70 }
71}
72
73#[derive(Debug, Clone)]
75struct IndexedFile {
76 path: PathBuf,
77 content: String,
78 size: usize,
79}
80
81pub struct FileSystemContextProvider {
83 config: FileSystemContextConfig,
84 cache: Arc<RwLock<HashMap<PathBuf, IndexedFile>>>,
86}
87
88impl FileSystemContextProvider {
89 pub fn new(config: FileSystemContextConfig) -> Self {
91 Self {
92 config,
93 cache: Arc::new(RwLock::new(HashMap::new())),
94 }
95 }
96
97 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 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 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}