mod chunking;mod embedding;mod embedding_index;mod indexing;mod project_index;mod project_index_debug_view;mod summary_backlog;mod summary_index;mod worktree_index; use anyhow::{Context as _, Result};use collections::HashMap;use fs::Fs;use gpui::{App, AppContext as _, AsyncApp, BorrowAppContext, Context, Entity, Global, WeakEntity};use language::LineEnding;use project::{Project, Worktree};use std::{    cmp::Ordering,    path::{Path, PathBuf},    sync::Arc,};use util::ResultExt as _;use workspace::Workspace; pub use embedding::*;pub use project_index::{LoadedSearchResult, ProjectIndex, SearchResult, Status};pub use project_index_debug_view::ProjectIndexDebugView;pub use summary_index::FileSummary; pub struct SemanticDb {    embedding_provider: Arc<dyn EmbeddingProvider>,    db_connection: Option<heed::Env>,    project_indices: HashMap<WeakEntity<Project>, Entity<ProjectIndex>>,} impl Global for SemanticDb {} impl SemanticDb {    pub async fn new(        db_path: PathBuf,        embedding_provider: Arc<dyn EmbeddingProvider>,        cx: &mut AsyncApp,    ) -> Result<Self> {        let db_connection = cx            .background_spawn(async move {                std::fs::create_dir_all(&db_path)?;                unsafe {                    heed::EnvOpenOptions::new()                        .map_size(1024 * 1024 * 1024)                        .max_dbs(3000)                        .open(db_path)                }            })            .await            .context("opening database connection")?;         cx.update(|cx| {            cx.observe_new(                |workspace: &mut Workspace, _window, cx: &mut Context<Workspace>| {                    let project = workspace.project().clone();                     if cx.has_global::<SemanticDb>() {                        cx.update_global::<SemanticDb, _>(|this, cx| {                            this.create_project_index(project, cx);                        })                    } else {                        log::info!("No SemanticDb, skipping project index")                    }                },            )            .detach();        })        .ok();         Ok(SemanticDb {            db_connection: Some(db_connection),            embedding_provider,            project_indices: HashMap::default(),        })    }     pub async fn load_results(        mut results: Vec<SearchResult>,        fs: &Arc<dyn Fs>,        cx: &AsyncApp,    ) -> Result<Vec<LoadedSearchResult>> {        let mut max_scores_by_path = HashMap::<_, (f32, usize)>::default();        for result in &results {            let (score, query_index) = max_scores_by_path                .entry((result.worktree.clone(), result.path.clone()))                .or_default();            if result.score > *score {                *score = result.score;                *query_index = result.query_index;            }        }         results.sort_by(|a, b| {            let max_score_a = max_scores_by_path[&(a.worktree.clone(), a.path.clone())].0;            let max_score_b = max_scores_by_path[&(b.worktree.clone(), b.path.clone())].0;            max_score_b                .partial_cmp(&max_score_a)                .unwrap_or(Ordering::Equal)                .then_with(|| a.worktree.entity_id().cmp(&b.worktree.entity_id()))                .then_with(|| a.path.cmp(&b.path))                .then_with(|| a.range.start.cmp(&b.range.start))        });         let mut last_loaded_file: Option<(Entity<Worktree>, Arc<Path>, PathBuf, String)> = None;        let mut loaded_results = Vec::<LoadedSearchResult>::new();        for result in results {            let full_path;            let file_content;            if let Some(last_loaded_file) =                last_loaded_file                    .as_ref()                    .filter(|(last_worktree, last_path, _, _)| {                        last_worktree == &result.worktree && last_path == &result.path                    })            {                full_path = last_loaded_file.2.clone();                file_content = &last_loaded_file.3;            } else {                let output = result.worktree.read_with(cx, |worktree, _cx| {                    let entry_abs_path = worktree.abs_path().join(&result.path);                    let mut entry_full_path = PathBuf::from(worktree.root_name());                    entry_full_path.push(&result.path);                    let file_content = async {                        let entry_abs_path = entry_abs_path;                        fs.load(&entry_abs_path).await                    };                    (entry_full_path, file_content)                })?;                full_path = output.0;                let Some(content) = output.1.await.log_err() else {                    continue;                };                last_loaded_file = Some((                    result.worktree.clone(),                    result.path.clone(),                    full_path.clone(),                    content,                ));                file_content = &last_loaded_file.as_ref().unwrap().3;            };             let query_index = max_scores_by_path[&(result.worktree.clone(), result.path.clone())].1;             let mut range_start = result.range.start.min(file_content.len());            let mut range_end = result.range.end.min(file_content.len());            while !file_content.is_char_boundary(range_start) {                range_start += 1;            }            while !file_content.is_char_boundary(range_end) {                range_end += 1;            }             let start_row = file_content[0..range_start].matches('\n').count() as u32;            let mut end_row = file_content[0..range_end].matches('\n').count() as u32;            let start_line_byte_offset = file_content[0..range_start]                .rfind('\n')                .map(|pos| pos + 1)                .unwrap_or_default();            let mut end_line_byte_offset = range_end;            if file_content[..end_line_byte_offset].ends_with('\n') {                end_row -= 1;            } else {                end_line_byte_offset = file_content[range_end..]                    .find('\n')                    .map(|pos| range_end + pos + 1)                    .unwrap_or_else(|| file_content.len());            }            let mut excerpt_content =                file_content[start_line_byte_offset..end_line_byte_offset].to_string();            LineEnding::normalize(&mut excerpt_content);             if let Some(prev_result) = loaded_results.last_mut() {                if prev_result.full_path == full_path {                    if *prev_result.row_range.end() + 1 == start_row {                        prev_result.row_range = *prev_result.row_range.start()..=end_row;                        prev_result.excerpt_content.push_str(&excerpt_content);                        continue;                    }                }            }             loaded_results.push(LoadedSearchResult {                path: result.path,                full_path,                excerpt_content,                row_range: start_row..=end_row,                query_index,            });        }         for result in &mut loaded_results {            while result.excerpt_content.ends_with("\n\n") {                result.excerpt_content.pop();                result.row_range =                    *result.row_range.start()..=result.row_range.end().saturating_sub(1)            }        }         Ok(loaded_results)    }     pub fn project_index(        &mut self,        project: Entity<Project>,        _cx: &mut App,    ) -> Option<Entity<ProjectIndex>> {        self.project_indices.get(&project.downgrade()).cloned()    }     pub fn remaining_summaries(        &self,        project: &WeakEntity<Project>,        cx: &mut App,    ) -> Option<usize> {        self.project_indices.get(project).map(|project_index| {            project_index.update(cx, |project_index, cx| {                project_index.remaining_summaries(cx)            })        })    }     pub fn create_project_index(        &mut self,        project: Entity<Project>,        cx: &mut App,    ) -> Entity<ProjectIndex> {        let project_index = cx.new(|cx| {            ProjectIndex::new(                project.clone(),                self.db_connection.clone().unwrap(),                self.embedding_provider.clone(),                cx,            )        });         let project_weak = project.downgrade();        self.project_indices            .insert(project_weak.clone(), project_index.clone());         cx.observe_release(&project, move |_, cx| {            if cx.has_global::<SemanticDb>() {                cx.update_global::<SemanticDb, _>(|this, _| {                    this.project_indices.remove(&project_weak);                })            }        })        .detach();         project_index    }} impl Drop for SemanticDb {    fn drop(&mut self) {        self.db_connection.take().unwrap().prepare_for_closing();    }} #[cfg(test)]mod tests {    use super::*;    use anyhow::anyhow;    use chunking::Chunk;    use embedding_index::{ChunkedFile, EmbeddingIndex};    use feature_flags::FeatureFlagAppExt;    use fs::FakeFs;    use futures::{future::BoxFuture, FutureExt};    use gpui::TestAppContext;    use indexing::IndexingEntrySet;    use language::language_settings::AllLanguageSettings;    use project::{Project, ProjectEntryId};    use serde_json::json;    use settings::SettingsStore;    use smol::channel;    use std::{future, path::Path, sync::Arc};    use util::separator;     fn init_test(cx: &mut TestAppContext) {        env_logger::try_init().ok();         cx.update(|cx| {            let store = SettingsStore::test(cx);            cx.set_global(store);            language::init(cx);            cx.update_flags(false, vec![]);            Project::init_settings(cx);            SettingsStore::update(cx, |store, cx| {                store.update_user_settings::<AllLanguageSettings>(cx, |_| {});            });        });    }     pub struct TestEmbeddingProvider {        batch_size: usize,        compute_embedding: Box<dyn Fn(&str) -> Result<Embedding> + Send + Sync>,    }     impl TestEmbeddingProvider {        pub fn new(            batch_size: usize,            compute_embedding: impl 'static + Fn(&str) -> Result<Embedding> + Send + Sync,        ) -> Self {            Self {                batch_size,                compute_embedding: Box::new(compute_embedding),            }        }    }     impl EmbeddingProvider for TestEmbeddingProvider {        fn embed<'a>(            &'a self,            texts: &'a [TextToEmbed<'a>],        ) -> BoxFuture<'a, Result<Vec<Embedding>>> {            let embeddings = texts                .iter()                .map(|to_embed| (self.compute_embedding)(to_embed.text))                .collect();            future::ready(embeddings).boxed()        }         fn batch_size(&self) -> usize {            self.batch_size        }    }     #[gpui::test]    async fn test_search(cx: &mut TestAppContext) {        cx.executor().allow_parking();         init_test(cx);         cx.update(|cx| {            // This functionality is staff-flagged.            cx.update_flags(true, vec![]);        });         let temp_dir = tempfile::tempdir().unwrap();         let mut semantic_index = SemanticDb::new(            temp_dir.path().into(),            Arc::new(TestEmbeddingProvider::new(16, |text| {                let mut embedding = vec![0f32; 2];                // if the text contains garbage, give it a 1 in the first dimension                if text.contains("garbage in") {                    embedding[0] = 0.9;                } else {                    embedding[0] = -0.9;                }                 if text.contains("garbage out") {                    embedding[1] = 0.9;                } else {                    embedding[1] = -0.9;                }                 Ok(Embedding::new(embedding))            })),            &mut cx.to_async(),        )        .await        .unwrap();         let fs = FakeFs::new(cx.executor());        let project_path = Path::new("/fake_project");         fs.insert_tree(            project_path,            json!({                "fixture": {                    "main.rs": include_str!("../fixture/main.rs"),                    "needle.md": include_str!("../fixture/needle.md"),                }            }),        )        .await;         let project = Project::test(fs, [project_path], cx).await;         let project_index = cx.update(|cx| {            let language_registry = project.read(cx).languages().clone();            let node_runtime = project.read(cx).node_runtime().unwrap().clone();            languages::init(language_registry, node_runtime, cx);            semantic_index.create_project_index(project.clone(), cx)        });         cx.run_until_parked();        while cx            .update(|cx| semantic_index.remaining_summaries(&project.downgrade(), cx))            .unwrap()            > 0        {            cx.run_until_parked();        }         let results = cx            .update(|cx| {                let project_index = project_index.read(cx);                let query = "garbage in, garbage out";                project_index.search(vec![query.into()], 4, cx)            })            .await            .unwrap();         assert!(            results.len() > 1,            "should have found some results, but only found {:?}",            results        );         for result in &results {            println!("result: {:?}", result.path);            println!("score: {:?}", result.score);        }         // Find result that is greater than 0.5        let search_result = results.iter().find(|result| result.score > 0.9).unwrap();         assert_eq!(            search_result.path.to_string_lossy(),            separator!("fixture/needle.md")        );         let content = cx            .update(|cx| {                let worktree = search_result.worktree.read(cx);                let entry_abs_path = worktree.abs_path().join(&search_result.path);                let fs = project.read(cx).fs().clone();                cx.background_spawn(async move { fs.load(&entry_abs_path).await.unwrap() })            })            .await;         let range = search_result.range.clone();        let content = content[range.clone()].to_owned();         assert!(content.contains("garbage in, garbage out"));    }     #[gpui::test]    async fn test_embed_files(cx: &mut TestAppContext) {        cx.executor().allow_parking();         let provider = Arc::new(TestEmbeddingProvider::new(3, |text| {            if text.contains('g') {                Err(anyhow!("cannot embed text containing a 'g' character"))            } else {                Ok(Embedding::new(                    ('a'..='z')                        .map(|char| text.chars().filter(|c| *c == char).count() as f32)                        .collect(),                ))            }        }));         let (indexing_progress_tx, _) = channel::unbounded();        let indexing_entries = Arc::new(IndexingEntrySet::new(indexing_progress_tx));         let (chunked_files_tx, chunked_files_rx) = channel::unbounded::<ChunkedFile>();        chunked_files_tx            .send_blocking(ChunkedFile {                path: Path::new("test1.md").into(),                mtime: None,                handle: indexing_entries.insert(ProjectEntryId::from_proto(0)),                text: "abcdefghijklmnop".to_string(),                chunks: [0..4, 4..8, 8..12, 12..16]                    .into_iter()                    .map(|range| Chunk {                        range,                        digest: Default::default(),                    })                    .collect(),            })            .unwrap();        chunked_files_tx            .send_blocking(ChunkedFile {                path: Path::new("test2.md").into(),                mtime: None,                handle: indexing_entries.insert(ProjectEntryId::from_proto(1)),                text: "qrstuvwxyz".to_string(),                chunks: [0..4, 4..8, 8..10]                    .into_iter()                    .map(|range| Chunk {                        range,                        digest: Default::default(),                    })                    .collect(),            })            .unwrap();        chunked_files_tx.close();         let embed_files_task =            cx.update(|cx| EmbeddingIndex::embed_files(provider.clone(), chunked_files_rx, cx));        embed_files_task.task.await.unwrap();         let embedded_files_rx = embed_files_task.files;

Lines 1–500 of 634. Load more