use log::info;
use std::path::PathBuf;
use crate::{
embedding::EmbeddingProvider,
error::{IndexerError, Result},
storage::{QdrantStore, SqliteStore},
};
pub struct SearchEngine {
#[allow(dead_code)]
sqlite_store: SqliteStore,
#[allow(dead_code)]
vector_store: QdrantStore,
#[allow(dead_code)]
embedding_provider: Box<dyn EmbeddingProvider>,
}
#[derive(Debug, Clone)]
pub struct SearchQuery {
pub text: String,
pub directory_filter: Option<PathBuf>,
pub limit: usize,
pub similarity_threshold: Option<f32>,
}
#[derive(Debug, Clone)]
pub struct SearchResult {
pub file_path: String,
pub chunk_id: usize,
pub score: f32,
pub parent_directories: Vec<String>,
}
impl SearchEngine {
pub fn new(
sqlite_store: SqliteStore,
vector_store: QdrantStore,
embedding_provider: Box<dyn EmbeddingProvider>,
) -> Self {
Self {
sqlite_store,
vector_store,
embedding_provider,
}
}
pub fn validate_query(&self, query: &SearchQuery) -> Result<()> {
Self::validate_query_static(query)
}
pub fn filter_results_by_directory(
&self,
results: Vec<SearchResult>,
directory_filter: &Option<PathBuf>,
) -> Vec<SearchResult> {
Self::filter_results_by_directory_static(results, directory_filter)
}
pub fn apply_similarity_threshold(
&self,
results: Vec<SearchResult>,
threshold: Option<f32>,
) -> Vec<SearchResult> {
Self::apply_similarity_threshold_static(results, threshold)
}
pub fn rank_results(&self, results: Vec<SearchResult>) -> Vec<SearchResult> {
Self::rank_results_static(results)
}
pub fn limit_results(&self, results: Vec<SearchResult>, limit: usize) -> Vec<SearchResult> {
Self::limit_results_static(results, limit)
}
pub fn validate_query_static(query: &SearchQuery) -> Result<()> {
if query.text.trim().is_empty() {
return Err(IndexerError::invalid_input("Search query cannot be empty"));
}
if query.limit == 0 {
return Err(IndexerError::invalid_input(
"Search limit must be greater than 0",
));
}
if let Some(threshold) = query.similarity_threshold {
if !(0.0..=1.0).contains(&threshold) {
return Err(IndexerError::invalid_input(
"Similarity threshold must be between 0.0 and 1.0",
));
}
}
if let Some(ref dir_filter) = query.directory_filter {
if !dir_filter.is_dir() && !dir_filter.exists() {
return Err(IndexerError::invalid_input(
"Directory filter must be a valid directory path",
));
}
}
Ok(())
}
pub fn filter_results_by_directory_static(
results: Vec<SearchResult>,
directory_filter: &Option<PathBuf>,
) -> Vec<SearchResult> {
if let Some(filter_dir) = directory_filter {
let filter_str = filter_dir.to_string_lossy();
results
.into_iter()
.filter(|result| result.file_path.starts_with(filter_str.as_ref()))
.collect()
} else {
results
}
}
pub fn apply_similarity_threshold_static(
results: Vec<SearchResult>,
threshold: Option<f32>,
) -> Vec<SearchResult> {
if let Some(min_score) = threshold {
results
.into_iter()
.filter(|result| result.score >= min_score)
.collect()
} else {
results
}
}
pub fn rank_results_static(mut results: Vec<SearchResult>) -> Vec<SearchResult> {
results.sort_by(|a, b| {
b.score
.partial_cmp(&a.score)
.unwrap_or(std::cmp::Ordering::Equal)
});
results
}
pub fn limit_results_static(results: Vec<SearchResult>, limit: usize) -> Vec<SearchResult> {
results.into_iter().take(limit).collect()
}
pub async fn search(&self, query: SearchQuery) -> Result<Vec<SearchResult>> {
let text = &query.text;
let limit = query.limit;
info!("Searching for: '{text}' with limit: {limit}");
self.validate_query(&query)?;
let query_embedding = self
.embedding_provider
.generate_embedding(query.text.clone())
.await?;
let search_results = self.vector_store.search(query_embedding, limit).await?;
let filtered_results =
self.filter_results_by_directory(search_results, &query.directory_filter);
let threshold_results =
self.apply_similarity_threshold(filtered_results, query.similarity_threshold);
let ranked_results = self.rank_results(threshold_results);
let final_results = self.limit_results(ranked_results, limit);
Ok(final_results)
}
pub async fn find_similar_files(
&self,
file_path: PathBuf,
limit: usize,
) -> Result<Vec<SearchResult>> {
info!("Finding files similar to: {file_path:?} with limit: {limit}");
if !file_path.exists() {
return Err(IndexerError::not_found(format!(
"File not found: {}",
file_path.display()
)));
}
if !file_path.is_file() {
return Err(IndexerError::invalid_input(format!(
"Path is not a file: {}",
file_path.display()
)));
}
let normalized_path = crate::utils::normalize_path(&file_path)?;
let file_record = self.sqlite_store.get_file_by_path(&normalized_path)?;
let file_embedding = if let Some(file_record) = file_record {
let chunks = match file_record.chunks_json {
Some(chunks_json) => {
serde_json::from_value::<Vec<String>>(chunks_json).map_err(|e| {
IndexerError::file_processing(format!("Failed to parse chunks: {e}"))
})?
}
None => {
return Err(IndexerError::not_found(format!(
"No chunks found for file: {}",
file_path.display()
)));
}
};
if chunks.is_empty() {
return Err(IndexerError::not_found(format!(
"No chunks found for file: {}",
file_path.display()
)));
}
let representative_chunk = &chunks[0];
self.embedding_provider
.generate_embedding(representative_chunk.clone())
.await?
} else {
let content = std::fs::read_to_string(&file_path)
.map_err(|e| IndexerError::file_processing(format!("Failed to read file: {e}")))?;
let representative_content = if content.len() > 512 {
&content[..512]
} else {
&content
};
self.embedding_provider
.generate_embedding(representative_content.to_string())
.await?
};
let search_results = self.vector_store.search(file_embedding, limit + 5).await?;
let mut file_scores: std::collections::HashMap<String, (f32, usize)> =
std::collections::HashMap::new();
let file_path_str = file_path.to_string_lossy().to_string();
for result in search_results {
if result.file_path == file_path_str {
continue;
}
let entry = file_scores
.entry(result.file_path.clone())
.or_insert((0.0, 0));
if result.score > entry.0 {
entry.0 = result.score;
entry.1 = result.chunk_id;
}
}
let mut similar_files: Vec<_> = file_scores.into_iter().collect();
similar_files.sort_by(|a, b| {
b.1 .0
.partial_cmp(&a.1 .0)
.unwrap_or(std::cmp::Ordering::Equal)
});
similar_files.truncate(limit);
let results: Vec<SearchResult> = similar_files
.into_iter()
.map(|(file_path, (score, chunk_id))| SearchResult {
file_path,
chunk_id,
score,
parent_directories: vec![], })
.collect();
Ok(results)
}
pub async fn get_file_content(
&self,
file_path: PathBuf,
chunk_range: Option<(usize, usize)>,
) -> Result<String> {
info!("Getting content for: {file_path:?} with chunks: {chunk_range:?}");
if !file_path.exists() {
return Err(IndexerError::not_found(format!(
"File not found: {}",
file_path.display()
)));
}
if !file_path.is_file() {
return Err(IndexerError::invalid_input(format!(
"Path is not a file: {}",
file_path.display()
)));
}
let normalized_path = crate::utils::normalize_path(&file_path)?;
let file_record = self.sqlite_store.get_file_by_path(&normalized_path)?;
let content = if let Some(file_record) = file_record {
if let Some(chunks_json) = file_record.chunks_json {
let chunks = serde_json::from_value::<Vec<String>>(chunks_json).map_err(|e| {
IndexerError::file_processing(format!("Failed to parse chunks: {e}"))
})?;
if let Some((start, end)) = chunk_range {
let start_idx = start.saturating_sub(1);
let end_idx = end.min(chunks.len());
if start_idx >= chunks.len() {
return Err(IndexerError::invalid_input(format!(
"Chunk range {start}-{end} exceeds available chunks ({})",
chunks.len()
)));
}
chunks[start_idx..end_idx].join("\n")
} else {
chunks.join("\n")
}
} else {
let content = std::fs::read_to_string(&file_path).map_err(|e| {
IndexerError::file_processing(format!("Failed to read file: {e}"))
})?;
if let Some((start, end)) = chunk_range {
let lines: Vec<&str> = content.lines().collect();
let lines_per_chunk = lines.len().div_ceil(10); let total_chunks = lines.len().div_ceil(lines_per_chunk);
if start > total_chunks || start == 0 {
return Err(IndexerError::invalid_input(format!(
"Chunk {start} is out of range. File has {total_chunks} estimated chunks"
)));
}
let start_line = (start - 1) * lines_per_chunk;
let end_line = (end * lines_per_chunk).min(lines.len());
lines[start_line..end_line].join("\n")
} else {
content
}
}
} else {
let content = std::fs::read_to_string(&file_path)
.map_err(|e| IndexerError::file_processing(format!("Failed to read file: {e}")))?;
if let Some((start, end)) = chunk_range {
let lines: Vec<&str> = content.lines().collect();
let lines_per_chunk = lines.len().div_ceil(10); let total_chunks = lines.len().div_ceil(lines_per_chunk);
if start > total_chunks || start == 0 {
return Err(IndexerError::invalid_input(format!(
"Chunk {start} is out of range. File has {total_chunks} estimated chunks"
)));
}
let start_line = (start - 1) * lines_per_chunk;
let end_line = (end * lines_per_chunk).min(lines.len());
lines[start_line..end_line].join("\n")
} else {
content
}
};
Ok(content)
}
}
pub async fn create_search_engine() -> Result<SearchEngine> {
let config = crate::Config::load()?;
crate::environment::validate_environment(&config).await?;
let sqlite_store = crate::storage::SqliteStore::new(&config.storage.sqlite_path)?;
let vector_store = crate::storage::QdrantStore::new(
&config.storage.qdrant.endpoint,
config.storage.qdrant.collection.clone(),
)
.await?;
let embedding_provider = crate::embedding::create_embedding_provider(&config.embedding)?;
Ok(SearchEngine::new(
sqlite_store,
vector_store,
embedding_provider,
))
}
#[cfg(test)]
mod tests {
use super::*;
use std::path::PathBuf;
fn create_sample_search_results() -> Vec<SearchResult> {
vec![
SearchResult {
file_path: "/home/user/docs/readme.md".to_string(),
chunk_id: 0,
score: 0.9,
parent_directories: vec!["docs".to_string()],
},
SearchResult {
file_path: "/home/user/code/main.rs".to_string(),
chunk_id: 1,
score: 0.8,
parent_directories: vec!["code".to_string()],
},
SearchResult {
file_path: "/home/user/docs/api.md".to_string(),
chunk_id: 0,
score: 0.7,
parent_directories: vec!["docs".to_string()],
},
SearchResult {
file_path: "/home/user/other/test.txt".to_string(),
chunk_id: 0,
score: 0.5,
parent_directories: vec!["other".to_string()],
},
]
}
#[test]
fn test_validate_query_success() {
let valid_query = SearchQuery {
text: "test search".to_string(),
directory_filter: None,
limit: 10,
similarity_threshold: Some(0.5),
};
assert!(SearchEngine::validate_query_static(&valid_query).is_ok());
}
#[test]
fn test_validate_query_empty_text() {
let invalid_query = SearchQuery {
text: "".to_string(),
directory_filter: None,
limit: 10,
similarity_threshold: None,
};
let result = SearchEngine::validate_query_static(&invalid_query);
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("cannot be empty"));
}
#[test]
fn test_validate_query_whitespace_only_text() {
let invalid_query = SearchQuery {
text: " \t\n ".to_string(),
directory_filter: None,
limit: 10,
similarity_threshold: None,
};
let result = SearchEngine::validate_query_static(&invalid_query);
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("cannot be empty"));
}
#[test]
fn test_validate_query_zero_limit() {
let invalid_query = SearchQuery {
text: "test".to_string(),
directory_filter: None,
limit: 0,
similarity_threshold: None,
};
let result = SearchEngine::validate_query_static(&invalid_query);
assert!(result.is_err());
assert!(result
.unwrap_err()
.to_string()
.contains("must be greater than 0"));
}
#[test]
fn test_validate_query_invalid_similarity_threshold() {
let invalid_queries = vec![
SearchQuery {
text: "test".to_string(),
directory_filter: None,
limit: 10,
similarity_threshold: Some(-0.1),
},
SearchQuery {
text: "test".to_string(),
directory_filter: None,
limit: 10,
similarity_threshold: Some(1.1),
},
];
for query in invalid_queries {
let result = SearchEngine::validate_query_static(&query);
assert!(result.is_err());
assert!(result
.unwrap_err()
.to_string()
.contains("between 0.0 and 1.0"));
}
}
#[test]
fn test_validate_query_valid_similarity_threshold() {
let valid_thresholds = vec![0.0, 0.5, 1.0];
for threshold in valid_thresholds {
let query = SearchQuery {
text: "test".to_string(),
directory_filter: None,
limit: 10,
similarity_threshold: Some(threshold),
};
assert!(SearchEngine::validate_query_static(&query).is_ok());
}
}
#[test]
fn test_filter_results_by_directory_with_filter() {
let results = create_sample_search_results();
let filter_dir = Some(PathBuf::from("/home/user/docs"));
let filtered = SearchEngine::filter_results_by_directory_static(results, &filter_dir);
assert_eq!(filtered.len(), 2);
assert!(filtered
.iter()
.all(|r| r.file_path.starts_with("/home/user/docs")));
}
#[test]
fn test_filter_results_by_directory_no_filter() {
let results = create_sample_search_results();
let original_count = results.len();
let filtered = SearchEngine::filter_results_by_directory_static(results, &None);
assert_eq!(filtered.len(), original_count);
}
#[test]
fn test_filter_results_by_directory_no_matches() {
let results = create_sample_search_results();
let filter_dir = Some(PathBuf::from("/nonexistent/path"));
let filtered = SearchEngine::filter_results_by_directory_static(results, &filter_dir);
assert_eq!(filtered.len(), 0);
}
#[test]
fn test_apply_similarity_threshold_with_threshold() {
let results = create_sample_search_results();
let threshold = Some(0.75);
let filtered = SearchEngine::apply_similarity_threshold_static(results, threshold);
assert_eq!(filtered.len(), 2);
assert!(filtered.iter().all(|r| r.score >= 0.75));
}
#[test]
fn test_apply_similarity_threshold_no_threshold() {
let results = create_sample_search_results();
let original_count = results.len();
let filtered = SearchEngine::apply_similarity_threshold_static(results, None);
assert_eq!(filtered.len(), original_count);
}
#[test]
fn test_apply_similarity_threshold_no_matches() {
let results = create_sample_search_results();
let threshold = Some(0.95);
let filtered = SearchEngine::apply_similarity_threshold_static(results, threshold);
assert_eq!(filtered.len(), 0);
}
#[test]
fn test_rank_results() {
let results = create_sample_search_results();
let ranked = SearchEngine::rank_results_static(results);
assert_eq!(ranked.len(), 4);
assert_eq!(ranked[0].score, 0.9);
assert_eq!(ranked[1].score, 0.8);
assert_eq!(ranked[2].score, 0.7);
assert_eq!(ranked[3].score, 0.5);
for i in 1..ranked.len() {
assert!(ranked[i - 1].score >= ranked[i].score);
}
}
#[test]
fn test_rank_results_empty() {
let ranked = SearchEngine::rank_results_static(vec![]);
assert_eq!(ranked.len(), 0);
}
#[test]
fn test_limit_results() {
let results = create_sample_search_results();
let limited = SearchEngine::limit_results_static(results, 2);
assert_eq!(limited.len(), 2);
}
#[test]
fn test_limit_results_larger_than_available() {
let results = create_sample_search_results();
let original_count = results.len();
let limited = SearchEngine::limit_results_static(results, 10);
assert_eq!(limited.len(), original_count);
}
#[test]
fn test_limit_results_zero() {
let results = create_sample_search_results();
let limited = SearchEngine::limit_results_static(results, 0);
assert_eq!(limited.len(), 0);
}
#[test]
fn test_search_query_creation() {
let query = SearchQuery {
text: "test query".to_string(),
directory_filter: Some(PathBuf::from("/test/dir")),
limit: 5,
similarity_threshold: Some(0.8),
};
assert_eq!(query.text, "test query");
assert_eq!(query.directory_filter, Some(PathBuf::from("/test/dir")));
assert_eq!(query.limit, 5);
assert_eq!(query.similarity_threshold, Some(0.8));
}
#[test]
fn test_search_result_creation() {
let result = SearchResult {
file_path: "/test/file.txt".to_string(),
chunk_id: 1,
score: 0.85,
parent_directories: vec!["test".to_string()],
};
assert_eq!(result.file_path, "/test/file.txt");
assert_eq!(result.chunk_id, 1);
assert_eq!(result.score, 0.85);
assert_eq!(result.parent_directories, vec!["test".to_string()]);
}
}