Semantic search v2 (#2789)
Move semantic search from navigation modal, to project search option. This PR is intended to be released in Preview only, and requires an opt-in semantic_index option to enable. Without this opt-in setting enable, the user should perceive no differences between previous project search. Release Notes: (Preview-only) - Added Semantic Search as a opt-in feature within Project Search - Show indexing feedback on indexing process within project search view
This commit is contained in:
Kyle Caverly
Generated
+770
-709
File diff suppressed because it is too large
Load Diff
+2
-2
@@ -63,7 +63,7 @@ members = [
|
||||
"crates/theme",
|
||||
"crates/theme_selector",
|
||||
"crates/util",
|
||||
"crates/vector_store",
|
||||
"crates/semantic_index",
|
||||
"crates/vim",
|
||||
"crates/vcs_menu",
|
||||
"crates/workspace",
|
||||
@@ -133,7 +133,7 @@ tree-sitter-yaml = { git = "https://github.com/zed-industries/tree-sitter-yaml",
|
||||
tree-sitter-lua = "0.0.14"
|
||||
|
||||
[patch.crates-io]
|
||||
tree-sitter = { git = "https://github.com/tree-sitter/tree-sitter", rev = "49226023693107fba9a1191136a4f47f38cdca73" }
|
||||
tree-sitter = { git = "https://github.com/tree-sitter/tree-sitter", rev = "1c65ca24bc9a734ab70115188f465e12eecf224e" }
|
||||
async-task = { git = "https://github.com/zed-industries/async-task", rev = "341b57d6de98cdfd7b418567b8de2022ca993a6e" }
|
||||
|
||||
# TODO - Remove when a version is released with this PR: https://github.com/servo/core-foundation-rs/pull/457
|
||||
|
||||
@@ -324,8 +324,8 @@
|
||||
// the terminal will default to matching the buffer's font family.
|
||||
// "font_family": "Zed Mono"
|
||||
},
|
||||
// Difference settings for vector_store
|
||||
"vector_store": {
|
||||
// Difference settings for semantic_index
|
||||
"semantic_index": {
|
||||
"enabled": false,
|
||||
"reindexing_delay_seconds": 600
|
||||
},
|
||||
|
||||
@@ -339,6 +339,8 @@ pub struct LanguageConfig {
|
||||
#[serde(default)]
|
||||
pub line_comment: Option<Arc<str>>,
|
||||
#[serde(default)]
|
||||
pub collapsed_placeholder: String,
|
||||
#[serde(default)]
|
||||
pub block_comment: Option<(Arc<str>, Arc<str>)>,
|
||||
#[serde(default)]
|
||||
pub overrides: HashMap<String, LanguageConfigOverride>,
|
||||
@@ -408,6 +410,7 @@ impl Default for LanguageConfig {
|
||||
line_comment: Default::default(),
|
||||
block_comment: Default::default(),
|
||||
overrides: Default::default(),
|
||||
collapsed_placeholder: Default::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -523,9 +526,10 @@ pub struct OutlineConfig {
|
||||
pub struct EmbeddingConfig {
|
||||
pub query: Query,
|
||||
pub item_capture_ix: u32,
|
||||
pub name_capture_ix: u32,
|
||||
pub name_capture_ix: Option<u32>,
|
||||
pub context_capture_ix: Option<u32>,
|
||||
pub extra_context_capture_ix: Option<u32>,
|
||||
pub collapse_capture_ix: Option<u32>,
|
||||
pub keep_capture_ix: Option<u32>,
|
||||
}
|
||||
|
||||
struct InjectionConfig {
|
||||
@@ -1247,23 +1251,26 @@ impl Language {
|
||||
let mut item_capture_ix = None;
|
||||
let mut name_capture_ix = None;
|
||||
let mut context_capture_ix = None;
|
||||
let mut extra_context_capture_ix = None;
|
||||
let mut collapse_capture_ix = None;
|
||||
let mut keep_capture_ix = None;
|
||||
get_capture_indices(
|
||||
&query,
|
||||
&mut [
|
||||
("item", &mut item_capture_ix),
|
||||
("name", &mut name_capture_ix),
|
||||
("context", &mut context_capture_ix),
|
||||
("context.extra", &mut extra_context_capture_ix),
|
||||
("keep", &mut keep_capture_ix),
|
||||
("collapse", &mut collapse_capture_ix),
|
||||
],
|
||||
);
|
||||
if let Some((item_capture_ix, name_capture_ix)) = item_capture_ix.zip(name_capture_ix) {
|
||||
if let Some(item_capture_ix) = item_capture_ix {
|
||||
grammar.embedding_config = Some(EmbeddingConfig {
|
||||
query,
|
||||
item_capture_ix,
|
||||
name_capture_ix,
|
||||
context_capture_ix,
|
||||
extra_context_capture_ix,
|
||||
collapse_capture_ix,
|
||||
keep_capture_ix,
|
||||
});
|
||||
}
|
||||
Ok(self)
|
||||
@@ -1548,9 +1555,20 @@ impl Language {
|
||||
pub fn grammar(&self) -> Option<&Arc<Grammar>> {
|
||||
self.grammar.as_ref()
|
||||
}
|
||||
|
||||
pub fn default_scope(self: &Arc<Self>) -> LanguageScope {
|
||||
LanguageScope {
|
||||
language: self.clone(),
|
||||
override_id: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl LanguageScope {
|
||||
pub fn collapsed_placeholder(&self) -> &str {
|
||||
self.language.config.collapsed_placeholder.as_ref()
|
||||
}
|
||||
|
||||
pub fn line_comment_prefix(&self) -> Option<&Arc<str>> {
|
||||
Override::as_option(
|
||||
self.config_override().map(|o| &o.line_comment),
|
||||
|
||||
@@ -20,6 +20,7 @@ settings = { path = "../settings" }
|
||||
theme = { path = "../theme" }
|
||||
util = { path = "../util" }
|
||||
workspace = { path = "../workspace" }
|
||||
semantic_index = { path = "../semantic_index" }
|
||||
anyhow.workspace = true
|
||||
futures.workspace = true
|
||||
log.workspace = true
|
||||
|
||||
@@ -18,7 +18,9 @@ use gpui::{
|
||||
Task, View, ViewContext, ViewHandle, WeakModelHandle, WeakViewHandle,
|
||||
};
|
||||
use menu::Confirm;
|
||||
use postage::stream::Stream;
|
||||
use project::{search::SearchQuery, Entry, Project};
|
||||
use semantic_index::SemanticIndex;
|
||||
use smallvec::SmallVec;
|
||||
use std::{
|
||||
any::{Any, TypeId},
|
||||
@@ -36,7 +38,10 @@ use workspace::{
|
||||
ItemNavHistory, Pane, ToolbarItemLocation, ToolbarItemView, Workspace, WorkspaceId,
|
||||
};
|
||||
|
||||
actions!(project_search, [SearchInNew, ToggleFocus, NextField]);
|
||||
actions!(
|
||||
project_search,
|
||||
[SearchInNew, ToggleFocus, NextField, ToggleSemanticSearch]
|
||||
);
|
||||
|
||||
#[derive(Default)]
|
||||
struct ActiveSearches(HashMap<WeakModelHandle<Project>, WeakViewHandle<ProjectSearchView>>);
|
||||
@@ -89,6 +94,7 @@ pub struct ProjectSearchView {
|
||||
model: ModelHandle<ProjectSearch>,
|
||||
query_editor: ViewHandle<Editor>,
|
||||
results_editor: ViewHandle<Editor>,
|
||||
semantic: Option<SemanticSearchState>,
|
||||
search_options: SearchOptions,
|
||||
panels_with_errors: HashSet<InputPanel>,
|
||||
active_match_index: Option<usize>,
|
||||
@@ -98,6 +104,12 @@ pub struct ProjectSearchView {
|
||||
excluded_files_editor: ViewHandle<Editor>,
|
||||
}
|
||||
|
||||
struct SemanticSearchState {
|
||||
file_count: usize,
|
||||
outstanding_file_count: usize,
|
||||
_progress_task: Task<()>,
|
||||
}
|
||||
|
||||
pub struct ProjectSearchBar {
|
||||
active_project_search: Option<ViewHandle<ProjectSearchView>>,
|
||||
subscription: Option<Subscription>,
|
||||
@@ -172,6 +184,63 @@ impl ProjectSearch {
|
||||
}));
|
||||
cx.notify();
|
||||
}
|
||||
|
||||
fn semantic_search(
|
||||
&mut self,
|
||||
query: String,
|
||||
include_files: Vec<GlobMatcher>,
|
||||
exclude_files: Vec<GlobMatcher>,
|
||||
cx: &mut ModelContext<Self>,
|
||||
) {
|
||||
let search = SemanticIndex::global(cx).map(|index| {
|
||||
index.update(cx, |semantic_index, cx| {
|
||||
semantic_index.search_project(
|
||||
self.project.clone(),
|
||||
query.clone(),
|
||||
10,
|
||||
include_files,
|
||||
exclude_files,
|
||||
cx,
|
||||
)
|
||||
})
|
||||
});
|
||||
self.search_id += 1;
|
||||
self.match_ranges.clear();
|
||||
self.pending_search = Some(cx.spawn(|this, mut cx| async move {
|
||||
let results = search?.await.log_err()?;
|
||||
|
||||
let (_task, mut match_ranges) = this.update(&mut cx, |this, cx| {
|
||||
this.excerpts.update(cx, |excerpts, cx| {
|
||||
excerpts.clear(cx);
|
||||
|
||||
let matches = results
|
||||
.into_iter()
|
||||
.map(|result| (result.buffer, vec![result.range.start..result.range.start]))
|
||||
.collect();
|
||||
|
||||
excerpts.stream_excerpts_with_context_lines(matches, 3, cx)
|
||||
})
|
||||
});
|
||||
|
||||
while let Some(match_range) = match_ranges.next().await {
|
||||
this.update(&mut cx, |this, cx| {
|
||||
this.match_ranges.push(match_range);
|
||||
while let Ok(Some(match_range)) = match_ranges.try_next() {
|
||||
this.match_ranges.push(match_range);
|
||||
}
|
||||
cx.notify();
|
||||
});
|
||||
}
|
||||
|
||||
this.update(&mut cx, |this, cx| {
|
||||
this.pending_search.take();
|
||||
cx.notify();
|
||||
});
|
||||
|
||||
None
|
||||
}));
|
||||
cx.notify();
|
||||
}
|
||||
}
|
||||
|
||||
pub enum ViewEvent {
|
||||
@@ -195,13 +264,24 @@ impl View for ProjectSearchView {
|
||||
enum Status {}
|
||||
|
||||
let theme = theme::current(cx).clone();
|
||||
let text = if self.query_editor.read(cx).text(cx).is_empty() {
|
||||
""
|
||||
} else if model.pending_search.is_some() {
|
||||
"Searching..."
|
||||
let text = if model.pending_search.is_some() {
|
||||
Cow::Borrowed("Searching...")
|
||||
} else if let Some(semantic) = &self.semantic {
|
||||
if semantic.outstanding_file_count > 0 {
|
||||
Cow::Owned(format!(
|
||||
"Indexing. {} of {}...",
|
||||
semantic.file_count - semantic.outstanding_file_count,
|
||||
semantic.file_count
|
||||
))
|
||||
} else {
|
||||
Cow::Borrowed("Indexing complete")
|
||||
}
|
||||
} else if self.query_editor.read(cx).text(cx).is_empty() {
|
||||
Cow::Borrowed("")
|
||||
} else {
|
||||
"No results"
|
||||
Cow::Borrowed("No results")
|
||||
};
|
||||
|
||||
MouseEventHandler::<Status, _>::new(0, cx, |_, _| {
|
||||
Label::new(text, theme.search.results_status.clone())
|
||||
.aligned()
|
||||
@@ -490,6 +570,7 @@ impl ProjectSearchView {
|
||||
model,
|
||||
query_editor,
|
||||
results_editor,
|
||||
semantic: None,
|
||||
search_options: options,
|
||||
panels_with_errors: HashSet::new(),
|
||||
active_match_index: None,
|
||||
@@ -577,11 +658,59 @@ impl ProjectSearchView {
|
||||
}
|
||||
|
||||
fn search(&mut self, cx: &mut ViewContext<Self>) {
|
||||
if let Some(semantic) = &mut self.semantic {
|
||||
if semantic.outstanding_file_count > 0 {
|
||||
return;
|
||||
}
|
||||
|
||||
let query = self.query_editor.read(cx).text(cx);
|
||||
if let Some((included_files, exclude_files)) =
|
||||
self.get_included_and_excluded_globsets(cx)
|
||||
{
|
||||
self.model.update(cx, |model, cx| {
|
||||
model.semantic_search(query, included_files, exclude_files, cx)
|
||||
});
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
if let Some(query) = self.build_search_query(cx) {
|
||||
self.model.update(cx, |model, cx| model.search(query, cx));
|
||||
}
|
||||
}
|
||||
|
||||
fn get_included_and_excluded_globsets(
|
||||
&mut self,
|
||||
cx: &mut ViewContext<Self>,
|
||||
) -> Option<(Vec<GlobMatcher>, Vec<GlobMatcher>)> {
|
||||
let included_files =
|
||||
match Self::load_glob_set(&self.included_files_editor.read(cx).text(cx)) {
|
||||
Ok(included_files) => {
|
||||
self.panels_with_errors.remove(&InputPanel::Include);
|
||||
included_files
|
||||
}
|
||||
Err(_e) => {
|
||||
self.panels_with_errors.insert(InputPanel::Include);
|
||||
cx.notify();
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let excluded_files =
|
||||
match Self::load_glob_set(&self.excluded_files_editor.read(cx).text(cx)) {
|
||||
Ok(excluded_files) => {
|
||||
self.panels_with_errors.remove(&InputPanel::Exclude);
|
||||
excluded_files
|
||||
}
|
||||
Err(_e) => {
|
||||
self.panels_with_errors.insert(InputPanel::Exclude);
|
||||
cx.notify();
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
Some((included_files, excluded_files))
|
||||
}
|
||||
|
||||
fn build_search_query(&mut self, cx: &mut ViewContext<Self>) -> Option<SearchQuery> {
|
||||
let text = self.query_editor.read(cx).text(cx);
|
||||
let included_files =
|
||||
@@ -873,6 +1002,7 @@ impl ProjectSearchBar {
|
||||
if let Some(search_view) = self.active_project_search.as_ref() {
|
||||
search_view.update(cx, |search_view, cx| {
|
||||
search_view.search_options.toggle(option);
|
||||
search_view.semantic = None;
|
||||
search_view.search(cx);
|
||||
});
|
||||
cx.notify();
|
||||
@@ -882,6 +1012,61 @@ impl ProjectSearchBar {
|
||||
}
|
||||
}
|
||||
|
||||
fn toggle_semantic_search(&mut self, cx: &mut ViewContext<Self>) -> bool {
|
||||
if let Some(search_view) = self.active_project_search.as_ref() {
|
||||
search_view.update(cx, |search_view, cx| {
|
||||
if search_view.semantic.is_some() {
|
||||
search_view.semantic = None;
|
||||
} else if let Some(semantic_index) = SemanticIndex::global(cx) {
|
||||
// TODO: confirm that it's ok to send this project
|
||||
search_view.search_options = SearchOptions::none();
|
||||
|
||||
let project = search_view.model.read(cx).project.clone();
|
||||
let index_task = semantic_index.update(cx, |semantic_index, cx| {
|
||||
semantic_index.index_project(project, cx)
|
||||
});
|
||||
|
||||
cx.spawn(|search_view, mut cx| async move {
|
||||
let (files_to_index, mut files_remaining_rx) = index_task.await?;
|
||||
|
||||
search_view.update(&mut cx, |search_view, cx| {
|
||||
cx.notify();
|
||||
search_view.semantic = Some(SemanticSearchState {
|
||||
file_count: files_to_index,
|
||||
outstanding_file_count: files_to_index,
|
||||
_progress_task: cx.spawn(|search_view, mut cx| async move {
|
||||
while let Some(count) = files_remaining_rx.recv().await {
|
||||
search_view
|
||||
.update(&mut cx, |search_view, cx| {
|
||||
if let Some(semantic_search_state) =
|
||||
&mut search_view.semantic
|
||||
{
|
||||
semantic_search_state.outstanding_file_count =
|
||||
count;
|
||||
cx.notify();
|
||||
if count == 0 {
|
||||
return;
|
||||
}
|
||||
}
|
||||
})
|
||||
.ok();
|
||||
}
|
||||
}),
|
||||
});
|
||||
})?;
|
||||
anyhow::Ok(())
|
||||
})
|
||||
.detach_and_log_err(cx);
|
||||
}
|
||||
cx.notify();
|
||||
});
|
||||
cx.notify();
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
fn render_nav_button(
|
||||
&self,
|
||||
icon: &'static str,
|
||||
@@ -959,6 +1144,42 @@ impl ProjectSearchBar {
|
||||
.into_any()
|
||||
}
|
||||
|
||||
fn render_semantic_search_button(&self, cx: &mut ViewContext<Self>) -> AnyElement<Self> {
|
||||
let tooltip_style = theme::current(cx).tooltip.clone();
|
||||
let is_active = if let Some(search) = self.active_project_search.as_ref() {
|
||||
let search = search.read(cx);
|
||||
search.semantic.is_some()
|
||||
} else {
|
||||
false
|
||||
};
|
||||
|
||||
let region_id = 3;
|
||||
|
||||
MouseEventHandler::<Self, _>::new(region_id, cx, |state, cx| {
|
||||
let theme = theme::current(cx);
|
||||
let style = theme
|
||||
.search
|
||||
.option_button
|
||||
.in_state(is_active)
|
||||
.style_for(state);
|
||||
Label::new("Semantic", style.text.clone())
|
||||
.contained()
|
||||
.with_style(style.container)
|
||||
})
|
||||
.on_click(MouseButton::Left, move |_, this, cx| {
|
||||
this.toggle_semantic_search(cx);
|
||||
})
|
||||
.with_cursor_style(CursorStyle::PointingHand)
|
||||
.with_tooltip::<Self>(
|
||||
region_id,
|
||||
format!("Toggle Semantic Search"),
|
||||
Some(Box::new(ToggleSemanticSearch)),
|
||||
tooltip_style,
|
||||
cx,
|
||||
)
|
||||
.into_any()
|
||||
}
|
||||
|
||||
fn is_option_enabled(&self, option: SearchOptions, cx: &AppContext) -> bool {
|
||||
if let Some(search) = self.active_project_search.as_ref() {
|
||||
search.read(cx).search_options.contains(option)
|
||||
@@ -1048,8 +1269,14 @@ impl View for ProjectSearchBar {
|
||||
.with_child(self.render_nav_button(">", Direction::Next, cx))
|
||||
.aligned(),
|
||||
)
|
||||
.with_child(
|
||||
Flex::row()
|
||||
.with_child({
|
||||
let row = if SemanticIndex::enabled(cx) {
|
||||
Flex::row().with_child(self.render_semantic_search_button(cx))
|
||||
} else {
|
||||
Flex::row()
|
||||
};
|
||||
|
||||
let row = row
|
||||
.with_child(self.render_option_button(
|
||||
"Case",
|
||||
SearchOptions::CASE_SENSITIVE,
|
||||
@@ -1067,8 +1294,10 @@ impl View for ProjectSearchBar {
|
||||
))
|
||||
.contained()
|
||||
.with_style(theme.search.option_button_group)
|
||||
.aligned(),
|
||||
)
|
||||
.aligned();
|
||||
|
||||
row
|
||||
})
|
||||
.contained()
|
||||
.with_margin_bottom(row_spacing),
|
||||
)
|
||||
|
||||
@@ -53,6 +53,10 @@ impl SearchOptions {
|
||||
}
|
||||
}
|
||||
|
||||
pub fn none() -> SearchOptions {
|
||||
SearchOptions::NONE
|
||||
}
|
||||
|
||||
pub fn from_query(query: &SearchQuery) -> SearchOptions {
|
||||
let mut options = SearchOptions::NONE;
|
||||
options.set(SearchOptions::WHOLE_WORD, query.whole_word());
|
||||
|
||||
@@ -1,11 +1,11 @@
|
||||
[package]
|
||||
name = "vector_store"
|
||||
name = "semantic_index"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
publish = false
|
||||
|
||||
[lib]
|
||||
path = "src/vector_store.rs"
|
||||
path = "src/semantic_index.rs"
|
||||
doctest = false
|
||||
|
||||
[dependencies]
|
||||
@@ -20,6 +20,7 @@ editor = { path = "../editor" }
|
||||
rpc = { path = "../rpc" }
|
||||
settings = { path = "../settings" }
|
||||
anyhow.workspace = true
|
||||
postage.workspace = true
|
||||
futures.workspace = true
|
||||
smol.workspace = true
|
||||
rusqlite = { version = "0.27.0", features = ["blob", "array", "modern_sqlite"] }
|
||||
@@ -33,8 +34,10 @@ async-trait.workspace = true
|
||||
bincode = "1.3.3"
|
||||
matrixmultiply = "0.3.7"
|
||||
tiktoken-rs = "0.5.0"
|
||||
parking_lot.workspace = true
|
||||
rand.workspace = true
|
||||
schemars.workspace = true
|
||||
globset.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
gpui = { path = "../gpui", features = ["test-support"] }
|
||||
@@ -43,7 +46,17 @@ project = { path = "../project", features = ["test-support"] }
|
||||
rpc = { path = "../rpc", features = ["test-support"] }
|
||||
workspace = { path = "../workspace", features = ["test-support"] }
|
||||
settings = { path = "../settings", features = ["test-support"]}
|
||||
tree-sitter-rust = "*"
|
||||
|
||||
pretty_assertions.workspace = true
|
||||
rand.workspace = true
|
||||
unindent.workspace = true
|
||||
tempdir.workspace = true
|
||||
ctor.workspace = true
|
||||
env_logger.workspace = true
|
||||
|
||||
tree-sitter-typescript = "*"
|
||||
tree-sitter-json = "*"
|
||||
tree-sitter-rust = "*"
|
||||
tree-sitter-toml = "*"
|
||||
tree-sitter-cpp = "*"
|
||||
tree-sitter-elixir = "*"
|
||||
@@ -1,20 +1,21 @@
|
||||
use std::{
|
||||
cmp::Ordering,
|
||||
collections::HashMap,
|
||||
path::{Path, PathBuf},
|
||||
rc::Rc,
|
||||
time::SystemTime,
|
||||
};
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
|
||||
use crate::parsing::ParsedFile;
|
||||
use crate::VECTOR_STORE_VERSION;
|
||||
use crate::{parsing::Document, SEMANTIC_INDEX_VERSION};
|
||||
use anyhow::{anyhow, Context, Result};
|
||||
use globset::GlobMatcher;
|
||||
use project::Fs;
|
||||
use rpc::proto::Timestamp;
|
||||
use rusqlite::{
|
||||
params,
|
||||
types::{FromSql, FromSqlResult, ValueRef},
|
||||
};
|
||||
use std::{
|
||||
cmp::Ordering,
|
||||
collections::HashMap,
|
||||
ops::Range,
|
||||
path::{Path, PathBuf},
|
||||
rc::Rc,
|
||||
sync::Arc,
|
||||
time::SystemTime,
|
||||
};
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct FileRecord {
|
||||
@@ -42,48 +43,94 @@ pub struct VectorDatabase {
|
||||
}
|
||||
|
||||
impl VectorDatabase {
|
||||
pub fn new(path: String) -> Result<Self> {
|
||||
pub async fn new(fs: Arc<dyn Fs>, path: Arc<PathBuf>) -> Result<Self> {
|
||||
if let Some(db_directory) = path.parent() {
|
||||
fs.create_dir(db_directory).await?;
|
||||
}
|
||||
|
||||
let this = Self {
|
||||
db: rusqlite::Connection::open(path)?,
|
||||
db: rusqlite::Connection::open(path.as_path())?,
|
||||
};
|
||||
this.initialize_database()?;
|
||||
Ok(this)
|
||||
}
|
||||
|
||||
fn get_existing_version(&self) -> Result<i64> {
|
||||
let mut version_query = self
|
||||
.db
|
||||
.prepare("SELECT version from semantic_index_config")?;
|
||||
version_query
|
||||
.query_row([], |row| Ok(row.get::<_, i64>(0)?))
|
||||
.map_err(|err| anyhow!("version query failed: {err}"))
|
||||
}
|
||||
|
||||
fn initialize_database(&self) -> Result<()> {
|
||||
rusqlite::vtab::array::load_module(&self.db)?;
|
||||
|
||||
// This will create the database if it doesnt exist
|
||||
// Delete existing tables, if SEMANTIC_INDEX_VERSION is bumped
|
||||
if self
|
||||
.get_existing_version()
|
||||
.map_or(false, |version| version == SEMANTIC_INDEX_VERSION as i64)
|
||||
{
|
||||
log::trace!("vector database schema up to date");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
log::trace!("vector database schema out of date. updating...");
|
||||
self.db
|
||||
.execute("DROP TABLE IF EXISTS documents", [])
|
||||
.context("failed to drop 'documents' table")?;
|
||||
self.db
|
||||
.execute("DROP TABLE IF EXISTS files", [])
|
||||
.context("failed to drop 'files' table")?;
|
||||
self.db
|
||||
.execute("DROP TABLE IF EXISTS worktrees", [])
|
||||
.context("failed to drop 'worktrees' table")?;
|
||||
self.db
|
||||
.execute("DROP TABLE IF EXISTS semantic_index_config", [])
|
||||
.context("failed to drop 'semantic_index_config' table")?;
|
||||
|
||||
// Initialize Vector Databasing Tables
|
||||
self.db.execute(
|
||||
"CREATE TABLE IF NOT EXISTS worktrees (
|
||||
"CREATE TABLE semantic_index_config (
|
||||
version INTEGER NOT NULL
|
||||
)",
|
||||
[],
|
||||
)?;
|
||||
|
||||
self.db.execute(
|
||||
"INSERT INTO semantic_index_config (version) VALUES (?1)",
|
||||
params![SEMANTIC_INDEX_VERSION],
|
||||
)?;
|
||||
|
||||
self.db.execute(
|
||||
"CREATE TABLE worktrees (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
absolute_path VARCHAR NOT NULL
|
||||
);
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS worktrees_absolute_path ON worktrees (absolute_path);
|
||||
CREATE UNIQUE INDEX worktrees_absolute_path ON worktrees (absolute_path);
|
||||
",
|
||||
[],
|
||||
)?;
|
||||
|
||||
self.db.execute(
|
||||
"CREATE TABLE IF NOT EXISTS files (
|
||||
"CREATE TABLE files (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
worktree_id INTEGER NOT NULL,
|
||||
relative_path VARCHAR NOT NULL,
|
||||
mtime_seconds INTEGER NOT NULL,
|
||||
mtime_nanos INTEGER NOT NULL,
|
||||
vector_store_version INTEGER NOT NULL,
|
||||
FOREIGN KEY(worktree_id) REFERENCES worktrees(id) ON DELETE CASCADE
|
||||
)",
|
||||
[],
|
||||
)?;
|
||||
|
||||
self.db.execute(
|
||||
"CREATE TABLE IF NOT EXISTS documents (
|
||||
"CREATE TABLE documents (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
file_id INTEGER NOT NULL,
|
||||
offset INTEGER NOT NULL,
|
||||
start_byte INTEGER NOT NULL,
|
||||
end_byte INTEGER NOT NULL,
|
||||
name VARCHAR NOT NULL,
|
||||
embedding BLOB NOT NULL,
|
||||
FOREIGN KEY(file_id) REFERENCES files(id) ON DELETE CASCADE
|
||||
@@ -91,6 +138,7 @@ impl VectorDatabase {
|
||||
[],
|
||||
)?;
|
||||
|
||||
log::trace!("vector database initialized with updated schema.");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -102,43 +150,44 @@ impl VectorDatabase {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn insert_file(&self, worktree_id: i64, indexed_file: ParsedFile) -> Result<()> {
|
||||
pub fn insert_file(
|
||||
&self,
|
||||
worktree_id: i64,
|
||||
path: PathBuf,
|
||||
mtime: SystemTime,
|
||||
documents: Vec<Document>,
|
||||
) -> Result<()> {
|
||||
// Write to files table, and return generated id.
|
||||
self.db.execute(
|
||||
"
|
||||
DELETE FROM files WHERE worktree_id = ?1 AND relative_path = ?2;
|
||||
",
|
||||
params![worktree_id, indexed_file.path.to_str()],
|
||||
params![worktree_id, path.to_str()],
|
||||
)?;
|
||||
let mtime = Timestamp::from(indexed_file.mtime);
|
||||
let mtime = Timestamp::from(mtime);
|
||||
self.db.execute(
|
||||
"
|
||||
INSERT INTO files
|
||||
(worktree_id, relative_path, mtime_seconds, mtime_nanos, vector_store_version)
|
||||
(worktree_id, relative_path, mtime_seconds, mtime_nanos)
|
||||
VALUES
|
||||
(?1, ?2, $3, $4, $5);
|
||||
(?1, ?2, $3, $4);
|
||||
",
|
||||
params![
|
||||
worktree_id,
|
||||
indexed_file.path.to_str(),
|
||||
mtime.seconds,
|
||||
mtime.nanos,
|
||||
VECTOR_STORE_VERSION
|
||||
],
|
||||
params![worktree_id, path.to_str(), mtime.seconds, mtime.nanos],
|
||||
)?;
|
||||
|
||||
let file_id = self.db.last_insert_rowid();
|
||||
|
||||
// Currently inserting at approximately 3400 documents a second
|
||||
// I imagine we can speed this up with a bulk insert of some kind.
|
||||
for document in indexed_file.documents {
|
||||
for document in documents {
|
||||
let embedding_blob = bincode::serialize(&document.embedding)?;
|
||||
|
||||
self.db.execute(
|
||||
"INSERT INTO documents (file_id, offset, name, embedding) VALUES (?1, ?2, ?3, ?4)",
|
||||
"INSERT INTO documents (file_id, start_byte, end_byte, name, embedding) VALUES (?1, ?2, ?3, ?4, $5)",
|
||||
params![
|
||||
file_id,
|
||||
document.offset.to_string(),
|
||||
document.range.start.to_string(),
|
||||
document.range.end.to_string(),
|
||||
document.name,
|
||||
embedding_blob
|
||||
],
|
||||
@@ -148,6 +197,23 @@ impl VectorDatabase {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn worktree_previously_indexed(&self, worktree_root_path: &Path) -> Result<bool> {
|
||||
let mut worktree_query = self
|
||||
.db
|
||||
.prepare("SELECT id FROM worktrees WHERE absolute_path = ?1")?;
|
||||
let worktree_id = worktree_query
|
||||
.query_row(params![worktree_root_path.to_string_lossy()], |row| {
|
||||
Ok(row.get::<_, i64>(0)?)
|
||||
})
|
||||
.map_err(|err| anyhow!(err));
|
||||
|
||||
if worktree_id.is_ok() {
|
||||
return Ok(true);
|
||||
} else {
|
||||
return Ok(false);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn find_or_create_worktree(&self, worktree_root_path: &Path) -> Result<i64> {
|
||||
// Check that the absolute path doesnt exist
|
||||
let mut worktree_query = self
|
||||
@@ -204,19 +270,26 @@ impl VectorDatabase {
|
||||
worktree_ids: &[i64],
|
||||
query_embedding: &Vec<f32>,
|
||||
limit: usize,
|
||||
) -> Result<Vec<(i64, PathBuf, usize, String)>> {
|
||||
include_globs: Vec<GlobMatcher>,
|
||||
exclude_globs: Vec<GlobMatcher>,
|
||||
) -> Result<Vec<(i64, PathBuf, Range<usize>)>> {
|
||||
let mut results = Vec::<(i64, f32)>::with_capacity(limit + 1);
|
||||
self.for_each_document(&worktree_ids, |id, embedding| {
|
||||
let similarity = dot(&embedding, &query_embedding);
|
||||
let ix = match results
|
||||
.binary_search_by(|(_, s)| similarity.partial_cmp(&s).unwrap_or(Ordering::Equal))
|
||||
{
|
||||
Ok(ix) => ix,
|
||||
Err(ix) => ix,
|
||||
};
|
||||
results.insert(ix, (id, similarity));
|
||||
results.truncate(limit);
|
||||
})?;
|
||||
self.for_each_document(
|
||||
&worktree_ids,
|
||||
include_globs,
|
||||
exclude_globs,
|
||||
|id, embedding| {
|
||||
let similarity = dot(&embedding, &query_embedding);
|
||||
let ix = match results.binary_search_by(|(_, s)| {
|
||||
similarity.partial_cmp(&s).unwrap_or(Ordering::Equal)
|
||||
}) {
|
||||
Ok(ix) => ix,
|
||||
Err(ix) => ix,
|
||||
};
|
||||
results.insert(ix, (id, similarity));
|
||||
results.truncate(limit);
|
||||
},
|
||||
)?;
|
||||
|
||||
let ids = results.into_iter().map(|(id, _)| id).collect::<Vec<_>>();
|
||||
self.get_documents_by_ids(&ids)
|
||||
@@ -225,22 +298,51 @@ impl VectorDatabase {
|
||||
fn for_each_document(
|
||||
&self,
|
||||
worktree_ids: &[i64],
|
||||
include_globs: Vec<GlobMatcher>,
|
||||
exclude_globs: Vec<GlobMatcher>,
|
||||
mut f: impl FnMut(i64, Vec<f32>),
|
||||
) -> Result<()> {
|
||||
let mut file_query = self.db.prepare(
|
||||
"
|
||||
SELECT
|
||||
id, relative_path
|
||||
FROM
|
||||
files
|
||||
WHERE
|
||||
worktree_id IN rarray(?)
|
||||
",
|
||||
)?;
|
||||
|
||||
let mut file_ids = Vec::<i64>::new();
|
||||
let mut rows = file_query.query([ids_to_sql(worktree_ids)])?;
|
||||
while let Some(row) = rows.next()? {
|
||||
let file_id = row.get(0)?;
|
||||
let relative_path = row.get_ref(1)?.as_str()?;
|
||||
let included = include_globs.is_empty()
|
||||
|| include_globs
|
||||
.iter()
|
||||
.any(|glob| glob.is_match(relative_path));
|
||||
let excluded = exclude_globs
|
||||
.iter()
|
||||
.any(|glob| glob.is_match(relative_path));
|
||||
if included && !excluded {
|
||||
file_ids.push(file_id);
|
||||
}
|
||||
}
|
||||
|
||||
let mut query_statement = self.db.prepare(
|
||||
"
|
||||
SELECT
|
||||
documents.id, documents.embedding
|
||||
id, embedding
|
||||
FROM
|
||||
documents, files
|
||||
documents
|
||||
WHERE
|
||||
documents.file_id = files.id AND
|
||||
files.worktree_id IN rarray(?)
|
||||
file_id IN rarray(?)
|
||||
",
|
||||
)?;
|
||||
|
||||
query_statement
|
||||
.query_map(params![ids_to_sql(worktree_ids)], |row| {
|
||||
.query_map(params![ids_to_sql(&file_ids)], |row| {
|
||||
Ok((row.get(0)?, row.get::<_, Embedding>(1)?))
|
||||
})?
|
||||
.filter_map(|row| row.ok())
|
||||
@@ -248,11 +350,15 @@ impl VectorDatabase {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn get_documents_by_ids(&self, ids: &[i64]) -> Result<Vec<(i64, PathBuf, usize, String)>> {
|
||||
fn get_documents_by_ids(&self, ids: &[i64]) -> Result<Vec<(i64, PathBuf, Range<usize>)>> {
|
||||
let mut statement = self.db.prepare(
|
||||
"
|
||||
SELECT
|
||||
documents.id, files.worktree_id, files.relative_path, documents.offset, documents.name
|
||||
documents.id,
|
||||
files.worktree_id,
|
||||
files.relative_path,
|
||||
documents.start_byte,
|
||||
documents.end_byte
|
||||
FROM
|
||||
documents, files
|
||||
WHERE
|
||||
@@ -266,15 +372,14 @@ impl VectorDatabase {
|
||||
row.get::<_, i64>(0)?,
|
||||
row.get::<_, i64>(1)?,
|
||||
row.get::<_, String>(2)?.into(),
|
||||
row.get(3)?,
|
||||
row.get(4)?,
|
||||
row.get(3)?..row.get(4)?,
|
||||
))
|
||||
})?;
|
||||
|
||||
let mut values_by_id = HashMap::<i64, (i64, PathBuf, usize, String)>::default();
|
||||
let mut values_by_id = HashMap::<i64, (i64, PathBuf, Range<usize>)>::default();
|
||||
for row in result_iter {
|
||||
let (id, worktree_id, path, offset, name) = row?;
|
||||
values_by_id.insert(id, (worktree_id, path, offset, name));
|
||||
let (id, worktree_id, path, range) = row?;
|
||||
values_by_id.insert(id, (worktree_id, path, range));
|
||||
}
|
||||
|
||||
let mut results = Vec::with_capacity(ids.len());
|
||||
@@ -67,17 +67,16 @@ impl EmbeddingProvider for DummyEmbeddings {
|
||||
}
|
||||
}
|
||||
|
||||
const INPUT_LIMIT: usize = 8190;
|
||||
const OPENAI_INPUT_LIMIT: usize = 8190;
|
||||
|
||||
impl OpenAIEmbeddings {
|
||||
fn truncate(span: String) -> String {
|
||||
let mut tokens = OPENAI_BPE_TOKENIZER.encode_with_special_tokens(span.as_ref());
|
||||
if tokens.len() > INPUT_LIMIT {
|
||||
tokens.truncate(INPUT_LIMIT);
|
||||
if tokens.len() > OPENAI_INPUT_LIMIT {
|
||||
tokens.truncate(OPENAI_INPUT_LIMIT);
|
||||
let result = OPENAI_BPE_TOKENIZER.decode(tokens.clone());
|
||||
if result.is_ok() {
|
||||
let transformed = result.unwrap();
|
||||
// assert_ne!(transformed, span);
|
||||
return transformed;
|
||||
}
|
||||
}
|
||||
@@ -88,6 +87,7 @@ impl OpenAIEmbeddings {
|
||||
async fn send_request(&self, api_key: &str, spans: Vec<&str>) -> Result<Response<AsyncBody>> {
|
||||
let request = Request::post("https://api.openai.com/v1/embeddings")
|
||||
.redirect_policy(isahc::config::RedirectPolicy::Follow)
|
||||
.timeout(Duration::from_secs(4))
|
||||
.header("Content-Type", "application/json")
|
||||
.header("Authorization", format!("Bearer {}", api_key))
|
||||
.body(
|
||||
@@ -106,7 +106,7 @@ impl OpenAIEmbeddings {
|
||||
#[async_trait]
|
||||
impl EmbeddingProvider for OpenAIEmbeddings {
|
||||
async fn embed_batch(&self, spans: Vec<&str>) -> Result<Vec<Vec<f32>>> {
|
||||
const BACKOFF_SECONDS: [usize; 3] = [65, 180, 360];
|
||||
const BACKOFF_SECONDS: [usize; 3] = [45, 75, 125];
|
||||
const MAX_RETRIES: usize = 3;
|
||||
|
||||
let api_key = OPENAI_API_KEY
|
||||
@@ -114,6 +114,7 @@ impl EmbeddingProvider for OpenAIEmbeddings {
|
||||
.ok_or_else(|| anyhow!("no api key"))?;
|
||||
|
||||
let mut request_number = 0;
|
||||
let mut truncated = false;
|
||||
let mut response: Response<AsyncBody>;
|
||||
let mut spans: Vec<String> = spans.iter().map(|x| x.to_string()).collect();
|
||||
while request_number < MAX_RETRIES {
|
||||
@@ -132,14 +133,25 @@ impl EmbeddingProvider for OpenAIEmbeddings {
|
||||
match response.status() {
|
||||
StatusCode::TOO_MANY_REQUESTS => {
|
||||
let delay = Duration::from_secs(BACKOFF_SECONDS[request_number - 1] as u64);
|
||||
log::trace!(
|
||||
"open ai rate limiting, delaying request by {:?} seconds",
|
||||
delay.as_secs()
|
||||
);
|
||||
self.executor.timer(delay).await;
|
||||
}
|
||||
StatusCode::BAD_REQUEST => {
|
||||
log::info!("BAD REQUEST: {:?}", &response.status());
|
||||
// Don't worry about delaying bad request, as we can assume
|
||||
// we haven't been rate limited yet.
|
||||
for span in spans.iter_mut() {
|
||||
*span = Self::truncate(span.to_string());
|
||||
// Only truncate if it hasnt been truncated before
|
||||
if !truncated {
|
||||
for span in spans.iter_mut() {
|
||||
*span = Self::truncate(span.clone());
|
||||
}
|
||||
truncated = true;
|
||||
} else {
|
||||
// If failing once already truncated, log the error and break the loop
|
||||
let mut body = String::new();
|
||||
response.body_mut().read_to_string(&mut body).await?;
|
||||
log::trace!("open ai bad request: {:?} {:?}", &response.status(), body);
|
||||
break;
|
||||
}
|
||||
}
|
||||
StatusCode::OK => {
|
||||
@@ -147,7 +159,7 @@ impl EmbeddingProvider for OpenAIEmbeddings {
|
||||
response.body_mut().read_to_string(&mut body).await?;
|
||||
let response: OpenAIEmbeddingResponse = serde_json::from_str(&body)?;
|
||||
|
||||
log::info!(
|
||||
log::trace!(
|
||||
"openai embedding completed. tokens: {:?}",
|
||||
response.usage.total_tokens
|
||||
);
|
||||
@@ -0,0 +1,299 @@
|
||||
use anyhow::{anyhow, Ok, Result};
|
||||
use language::{Grammar, Language};
|
||||
use std::{
|
||||
cmp::{self, Reverse},
|
||||
collections::HashSet,
|
||||
ops::Range,
|
||||
path::Path,
|
||||
sync::Arc,
|
||||
};
|
||||
use tree_sitter::{Parser, QueryCursor};
|
||||
|
||||
#[derive(Debug, PartialEq, Clone)]
|
||||
pub struct Document {
|
||||
pub name: String,
|
||||
pub range: Range<usize>,
|
||||
pub content: String,
|
||||
pub embedding: Vec<f32>,
|
||||
}
|
||||
|
||||
const CODE_CONTEXT_TEMPLATE: &str =
|
||||
"The below code snippet is from file '<path>'\n\n```<language>\n<item>\n```";
|
||||
const ENTIRE_FILE_TEMPLATE: &str =
|
||||
"The below snippet is from file '<path>'\n\n```<language>\n<item>\n```";
|
||||
pub const PARSEABLE_ENTIRE_FILE_TYPES: &[&str] = &["TOML", "YAML", "CSS"];
|
||||
|
||||
pub struct CodeContextRetriever {
|
||||
pub parser: Parser,
|
||||
pub cursor: QueryCursor,
|
||||
}
|
||||
|
||||
// Every match has an item, this represents the fundamental treesitter symbol and anchors the search
|
||||
// Every match has one or more 'name' captures. These indicate the display range of the item for deduplication.
|
||||
// If there are preceeding comments, we track this with a context capture
|
||||
// If there is a piece that should be collapsed in hierarchical queries, we capture it with a collapse capture
|
||||
// If there is a piece that should be kept inside a collapsed node, we capture it with a keep capture
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CodeContextMatch {
|
||||
pub start_col: usize,
|
||||
pub item_range: Option<Range<usize>>,
|
||||
pub name_range: Option<Range<usize>>,
|
||||
pub context_ranges: Vec<Range<usize>>,
|
||||
pub collapse_ranges: Vec<Range<usize>>,
|
||||
}
|
||||
|
||||
impl CodeContextRetriever {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
parser: Parser::new(),
|
||||
cursor: QueryCursor::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_entire_file(
|
||||
&self,
|
||||
relative_path: &Path,
|
||||
language_name: Arc<str>,
|
||||
content: &str,
|
||||
) -> Result<Vec<Document>> {
|
||||
let document_span = ENTIRE_FILE_TEMPLATE
|
||||
.replace("<path>", relative_path.to_string_lossy().as_ref())
|
||||
.replace("<language>", language_name.as_ref())
|
||||
.replace("item", &content);
|
||||
|
||||
Ok(vec![Document {
|
||||
range: 0..content.len(),
|
||||
content: document_span,
|
||||
embedding: Vec::new(),
|
||||
name: language_name.to_string(),
|
||||
}])
|
||||
}
|
||||
|
||||
fn get_matches_in_file(
|
||||
&mut self,
|
||||
content: &str,
|
||||
grammar: &Arc<Grammar>,
|
||||
) -> Result<Vec<CodeContextMatch>> {
|
||||
let embedding_config = grammar
|
||||
.embedding_config
|
||||
.as_ref()
|
||||
.ok_or_else(|| anyhow!("no embedding queries"))?;
|
||||
self.parser.set_language(grammar.ts_language).unwrap();
|
||||
|
||||
let tree = self
|
||||
.parser
|
||||
.parse(&content, None)
|
||||
.ok_or_else(|| anyhow!("parsing failed"))?;
|
||||
|
||||
let mut captures: Vec<CodeContextMatch> = Vec::new();
|
||||
let mut collapse_ranges: Vec<Range<usize>> = Vec::new();
|
||||
let mut keep_ranges: Vec<Range<usize>> = Vec::new();
|
||||
for mat in self.cursor.matches(
|
||||
&embedding_config.query,
|
||||
tree.root_node(),
|
||||
content.as_bytes(),
|
||||
) {
|
||||
let mut start_col = 0;
|
||||
let mut item_range: Option<Range<usize>> = None;
|
||||
let mut name_range: Option<Range<usize>> = None;
|
||||
let mut context_ranges: Vec<Range<usize>> = Vec::new();
|
||||
collapse_ranges.clear();
|
||||
keep_ranges.clear();
|
||||
for capture in mat.captures {
|
||||
if capture.index == embedding_config.item_capture_ix {
|
||||
item_range = Some(capture.node.byte_range());
|
||||
start_col = capture.node.start_position().column;
|
||||
} else if Some(capture.index) == embedding_config.name_capture_ix {
|
||||
name_range = Some(capture.node.byte_range());
|
||||
} else if Some(capture.index) == embedding_config.context_capture_ix {
|
||||
context_ranges.push(capture.node.byte_range());
|
||||
} else if Some(capture.index) == embedding_config.collapse_capture_ix {
|
||||
collapse_ranges.push(capture.node.byte_range());
|
||||
} else if Some(capture.index) == embedding_config.keep_capture_ix {
|
||||
keep_ranges.push(capture.node.byte_range());
|
||||
}
|
||||
}
|
||||
|
||||
captures.push(CodeContextMatch {
|
||||
start_col,
|
||||
item_range,
|
||||
name_range,
|
||||
context_ranges,
|
||||
collapse_ranges: subtract_ranges(&collapse_ranges, &keep_ranges),
|
||||
});
|
||||
}
|
||||
Ok(captures)
|
||||
}
|
||||
|
||||
pub fn parse_file_with_template(
|
||||
&mut self,
|
||||
relative_path: &Path,
|
||||
content: &str,
|
||||
language: Arc<Language>,
|
||||
) -> Result<Vec<Document>> {
|
||||
let language_name = language.name();
|
||||
|
||||
if PARSEABLE_ENTIRE_FILE_TYPES.contains(&language_name.as_ref()) {
|
||||
return self.parse_entire_file(relative_path, language_name, &content);
|
||||
}
|
||||
|
||||
let mut documents = self.parse_file(content, language)?;
|
||||
for document in &mut documents {
|
||||
document.content = CODE_CONTEXT_TEMPLATE
|
||||
.replace("<path>", relative_path.to_string_lossy().as_ref())
|
||||
.replace("<language>", language_name.as_ref())
|
||||
.replace("item", &document.content);
|
||||
}
|
||||
Ok(documents)
|
||||
}
|
||||
|
||||
pub fn parse_file(&mut self, content: &str, language: Arc<Language>) -> Result<Vec<Document>> {
|
||||
let grammar = language
|
||||
.grammar()
|
||||
.ok_or_else(|| anyhow!("no grammar for language"))?;
|
||||
|
||||
// Iterate through query matches
|
||||
let matches = self.get_matches_in_file(content, grammar)?;
|
||||
|
||||
let language_scope = language.default_scope();
|
||||
let placeholder = language_scope.collapsed_placeholder();
|
||||
|
||||
let mut documents = Vec::new();
|
||||
let mut collapsed_ranges_within = Vec::new();
|
||||
let mut parsed_name_ranges = HashSet::new();
|
||||
for (i, context_match) in matches.iter().enumerate() {
|
||||
// Items which are collapsible but not embeddable have no item range
|
||||
let item_range = if let Some(item_range) = context_match.item_range.clone() {
|
||||
item_range
|
||||
} else {
|
||||
continue;
|
||||
};
|
||||
|
||||
// Checks for deduplication
|
||||
let name;
|
||||
if let Some(name_range) = context_match.name_range.clone() {
|
||||
name = content
|
||||
.get(name_range.clone())
|
||||
.map_or(String::new(), |s| s.to_string());
|
||||
if parsed_name_ranges.contains(&name_range) {
|
||||
continue;
|
||||
}
|
||||
parsed_name_ranges.insert(name_range);
|
||||
} else {
|
||||
name = String::new();
|
||||
}
|
||||
|
||||
collapsed_ranges_within.clear();
|
||||
'outer: for remaining_match in &matches[(i + 1)..] {
|
||||
for collapsed_range in &remaining_match.collapse_ranges {
|
||||
if item_range.start <= collapsed_range.start
|
||||
&& item_range.end >= collapsed_range.end
|
||||
{
|
||||
collapsed_ranges_within.push(collapsed_range.clone());
|
||||
} else {
|
||||
break 'outer;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
collapsed_ranges_within.sort_by_key(|r| (r.start, Reverse(r.end)));
|
||||
|
||||
let mut document_content = String::new();
|
||||
for context_range in &context_match.context_ranges {
|
||||
document_content.push_str(&content[context_range.clone()]);
|
||||
document_content.push_str("\n");
|
||||
}
|
||||
|
||||
let mut offset = item_range.start;
|
||||
for collapsed_range in &collapsed_ranges_within {
|
||||
if collapsed_range.start > offset {
|
||||
add_content_from_range(
|
||||
&mut document_content,
|
||||
content,
|
||||
offset..collapsed_range.start,
|
||||
context_match.start_col,
|
||||
);
|
||||
offset = collapsed_range.start;
|
||||
}
|
||||
|
||||
if collapsed_range.end > offset {
|
||||
document_content.push_str(placeholder);
|
||||
offset = collapsed_range.end;
|
||||
}
|
||||
}
|
||||
|
||||
if offset < item_range.end {
|
||||
add_content_from_range(
|
||||
&mut document_content,
|
||||
content,
|
||||
offset..item_range.end,
|
||||
context_match.start_col,
|
||||
);
|
||||
}
|
||||
|
||||
documents.push(Document {
|
||||
name,
|
||||
content: document_content,
|
||||
range: item_range.clone(),
|
||||
embedding: vec![],
|
||||
})
|
||||
}
|
||||
|
||||
return Ok(documents);
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn subtract_ranges(
|
||||
ranges: &[Range<usize>],
|
||||
ranges_to_subtract: &[Range<usize>],
|
||||
) -> Vec<Range<usize>> {
|
||||
let mut result = Vec::new();
|
||||
|
||||
let mut ranges_to_subtract = ranges_to_subtract.iter().peekable();
|
||||
|
||||
for range in ranges {
|
||||
let mut offset = range.start;
|
||||
|
||||
while offset < range.end {
|
||||
if let Some(range_to_subtract) = ranges_to_subtract.peek() {
|
||||
if offset < range_to_subtract.start {
|
||||
let next_offset = cmp::min(range_to_subtract.start, range.end);
|
||||
result.push(offset..next_offset);
|
||||
offset = next_offset;
|
||||
} else {
|
||||
let next_offset = cmp::min(range_to_subtract.end, range.end);
|
||||
offset = next_offset;
|
||||
}
|
||||
|
||||
if offset >= range_to_subtract.end {
|
||||
ranges_to_subtract.next();
|
||||
}
|
||||
} else {
|
||||
result.push(offset..range.end);
|
||||
offset = range.end;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
fn add_content_from_range(
|
||||
output: &mut String,
|
||||
content: &str,
|
||||
range: Range<usize>,
|
||||
start_col: usize,
|
||||
) {
|
||||
for mut line in content.get(range.clone()).unwrap_or("").lines() {
|
||||
for _ in 0..start_col {
|
||||
if line.starts_with(' ') {
|
||||
line = &line[1..];
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
output.push_str(line);
|
||||
output.push('\n');
|
||||
}
|
||||
output.pop();
|
||||
}
|
||||
@@ -0,0 +1,777 @@
|
||||
mod db;
|
||||
mod embedding;
|
||||
mod parsing;
|
||||
pub mod semantic_index_settings;
|
||||
|
||||
#[cfg(test)]
|
||||
mod semantic_index_tests;
|
||||
|
||||
use crate::semantic_index_settings::SemanticIndexSettings;
|
||||
use anyhow::{anyhow, Result};
|
||||
use db::VectorDatabase;
|
||||
use embedding::{EmbeddingProvider, OpenAIEmbeddings};
|
||||
use futures::{channel::oneshot, Future};
|
||||
use globset::GlobMatcher;
|
||||
use gpui::{AppContext, AsyncAppContext, Entity, ModelContext, ModelHandle, Task, WeakModelHandle};
|
||||
use language::{Anchor, Buffer, Language, LanguageRegistry};
|
||||
use parking_lot::Mutex;
|
||||
use parsing::{CodeContextRetriever, Document, PARSEABLE_ENTIRE_FILE_TYPES};
|
||||
use postage::watch;
|
||||
use project::{Fs, Project, WorktreeId};
|
||||
use smol::channel;
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
mem,
|
||||
ops::Range,
|
||||
path::{Path, PathBuf},
|
||||
sync::{Arc, Weak},
|
||||
time::{Instant, SystemTime},
|
||||
};
|
||||
use util::{
|
||||
channel::{ReleaseChannel, RELEASE_CHANNEL, RELEASE_CHANNEL_NAME},
|
||||
http::HttpClient,
|
||||
paths::EMBEDDINGS_DIR,
|
||||
ResultExt,
|
||||
};
|
||||
|
||||
const SEMANTIC_INDEX_VERSION: usize = 6;
|
||||
const EMBEDDINGS_BATCH_SIZE: usize = 80;
|
||||
|
||||
pub fn init(
|
||||
fs: Arc<dyn Fs>,
|
||||
http_client: Arc<dyn HttpClient>,
|
||||
language_registry: Arc<LanguageRegistry>,
|
||||
cx: &mut AppContext,
|
||||
) {
|
||||
settings::register::<SemanticIndexSettings>(cx);
|
||||
|
||||
let db_file_path = EMBEDDINGS_DIR
|
||||
.join(Path::new(RELEASE_CHANNEL_NAME.as_str()))
|
||||
.join("embeddings_db");
|
||||
|
||||
if *RELEASE_CHANNEL == ReleaseChannel::Stable
|
||||
|| !settings::get::<SemanticIndexSettings>(cx).enabled
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
cx.spawn(move |mut cx| async move {
|
||||
let semantic_index = SemanticIndex::new(
|
||||
fs,
|
||||
db_file_path,
|
||||
Arc::new(OpenAIEmbeddings {
|
||||
client: http_client,
|
||||
executor: cx.background(),
|
||||
}),
|
||||
language_registry,
|
||||
cx.clone(),
|
||||
)
|
||||
.await?;
|
||||
|
||||
cx.update(|cx| {
|
||||
cx.set_global(semantic_index.clone());
|
||||
});
|
||||
|
||||
anyhow::Ok(())
|
||||
})
|
||||
.detach();
|
||||
}
|
||||
|
||||
pub struct SemanticIndex {
|
||||
fs: Arc<dyn Fs>,
|
||||
database_url: Arc<PathBuf>,
|
||||
embedding_provider: Arc<dyn EmbeddingProvider>,
|
||||
language_registry: Arc<LanguageRegistry>,
|
||||
db_update_tx: channel::Sender<DbOperation>,
|
||||
parsing_files_tx: channel::Sender<PendingFile>,
|
||||
_db_update_task: Task<()>,
|
||||
_embed_batch_tasks: Vec<Task<()>>,
|
||||
_batch_files_task: Task<()>,
|
||||
_parsing_files_tasks: Vec<Task<()>>,
|
||||
projects: HashMap<WeakModelHandle<Project>, ProjectState>,
|
||||
}
|
||||
|
||||
struct ProjectState {
|
||||
worktree_db_ids: Vec<(WorktreeId, i64)>,
|
||||
outstanding_job_count_rx: watch::Receiver<usize>,
|
||||
_outstanding_job_count_tx: Arc<Mutex<watch::Sender<usize>>>,
|
||||
}
|
||||
|
||||
struct JobHandle {
|
||||
tx: Weak<Mutex<watch::Sender<usize>>>,
|
||||
}
|
||||
|
||||
impl ProjectState {
|
||||
fn db_id_for_worktree_id(&self, id: WorktreeId) -> Option<i64> {
|
||||
self.worktree_db_ids
|
||||
.iter()
|
||||
.find_map(|(worktree_id, db_id)| {
|
||||
if *worktree_id == id {
|
||||
Some(*db_id)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn worktree_id_for_db_id(&self, id: i64) -> Option<WorktreeId> {
|
||||
self.worktree_db_ids
|
||||
.iter()
|
||||
.find_map(|(worktree_id, db_id)| {
|
||||
if *db_id == id {
|
||||
Some(*worktree_id)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub struct PendingFile {
|
||||
worktree_db_id: i64,
|
||||
relative_path: PathBuf,
|
||||
absolute_path: PathBuf,
|
||||
language: Arc<Language>,
|
||||
modified_time: SystemTime,
|
||||
job_handle: JobHandle,
|
||||
}
|
||||
|
||||
pub struct SearchResult {
|
||||
pub buffer: ModelHandle<Buffer>,
|
||||
pub range: Range<Anchor>,
|
||||
}
|
||||
|
||||
enum DbOperation {
|
||||
InsertFile {
|
||||
worktree_id: i64,
|
||||
documents: Vec<Document>,
|
||||
path: PathBuf,
|
||||
mtime: SystemTime,
|
||||
job_handle: JobHandle,
|
||||
},
|
||||
Delete {
|
||||
worktree_id: i64,
|
||||
path: PathBuf,
|
||||
},
|
||||
FindOrCreateWorktree {
|
||||
path: PathBuf,
|
||||
sender: oneshot::Sender<Result<i64>>,
|
||||
},
|
||||
FileMTimes {
|
||||
worktree_id: i64,
|
||||
sender: oneshot::Sender<Result<HashMap<PathBuf, SystemTime>>>,
|
||||
},
|
||||
WorktreePreviouslyIndexed {
|
||||
path: Arc<Path>,
|
||||
sender: oneshot::Sender<Result<bool>>,
|
||||
},
|
||||
}
|
||||
|
||||
enum EmbeddingJob {
|
||||
Enqueue {
|
||||
worktree_id: i64,
|
||||
path: PathBuf,
|
||||
mtime: SystemTime,
|
||||
documents: Vec<Document>,
|
||||
job_handle: JobHandle,
|
||||
},
|
||||
Flush,
|
||||
}
|
||||
|
||||
impl SemanticIndex {
|
||||
pub fn global(cx: &AppContext) -> Option<ModelHandle<SemanticIndex>> {
|
||||
if cx.has_global::<ModelHandle<Self>>() {
|
||||
Some(cx.global::<ModelHandle<SemanticIndex>>().clone())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
pub fn enabled(cx: &AppContext) -> bool {
|
||||
settings::get::<SemanticIndexSettings>(cx).enabled
|
||||
&& *RELEASE_CHANNEL != ReleaseChannel::Stable
|
||||
}
|
||||
|
||||
async fn new(
|
||||
fs: Arc<dyn Fs>,
|
||||
database_url: PathBuf,
|
||||
embedding_provider: Arc<dyn EmbeddingProvider>,
|
||||
language_registry: Arc<LanguageRegistry>,
|
||||
mut cx: AsyncAppContext,
|
||||
) -> Result<ModelHandle<Self>> {
|
||||
let t0 = Instant::now();
|
||||
let database_url = Arc::new(database_url);
|
||||
|
||||
let db = cx
|
||||
.background()
|
||||
.spawn(VectorDatabase::new(fs.clone(), database_url.clone()))
|
||||
.await?;
|
||||
|
||||
log::trace!(
|
||||
"db initialization took {:?} milliseconds",
|
||||
t0.elapsed().as_millis()
|
||||
);
|
||||
|
||||
Ok(cx.add_model(|cx| {
|
||||
let t0 = Instant::now();
|
||||
// Perform database operations
|
||||
let (db_update_tx, db_update_rx) = channel::unbounded();
|
||||
let _db_update_task = cx.background().spawn({
|
||||
async move {
|
||||
while let Ok(job) = db_update_rx.recv().await {
|
||||
Self::run_db_operation(&db, job)
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// Group documents into batches and send them to the embedding provider.
|
||||
let (embed_batch_tx, embed_batch_rx) =
|
||||
channel::unbounded::<Vec<(i64, Vec<Document>, PathBuf, SystemTime, JobHandle)>>();
|
||||
let mut _embed_batch_tasks = Vec::new();
|
||||
for _ in 0..cx.background().num_cpus() {
|
||||
let embed_batch_rx = embed_batch_rx.clone();
|
||||
_embed_batch_tasks.push(cx.background().spawn({
|
||||
let db_update_tx = db_update_tx.clone();
|
||||
let embedding_provider = embedding_provider.clone();
|
||||
async move {
|
||||
while let Ok(embeddings_queue) = embed_batch_rx.recv().await {
|
||||
Self::compute_embeddings_for_batch(
|
||||
embeddings_queue,
|
||||
&embedding_provider,
|
||||
&db_update_tx,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}));
|
||||
}
|
||||
|
||||
// Group documents into batches and send them to the embedding provider.
|
||||
let (batch_files_tx, batch_files_rx) = channel::unbounded::<EmbeddingJob>();
|
||||
let _batch_files_task = cx.background().spawn(async move {
|
||||
let mut queue_len = 0;
|
||||
let mut embeddings_queue = vec![];
|
||||
while let Ok(job) = batch_files_rx.recv().await {
|
||||
Self::enqueue_documents_to_embed(
|
||||
job,
|
||||
&mut queue_len,
|
||||
&mut embeddings_queue,
|
||||
&embed_batch_tx,
|
||||
);
|
||||
}
|
||||
});
|
||||
|
||||
// Parse files into embeddable documents.
|
||||
let (parsing_files_tx, parsing_files_rx) = channel::unbounded::<PendingFile>();
|
||||
let mut _parsing_files_tasks = Vec::new();
|
||||
for _ in 0..cx.background().num_cpus() {
|
||||
let fs = fs.clone();
|
||||
let parsing_files_rx = parsing_files_rx.clone();
|
||||
let batch_files_tx = batch_files_tx.clone();
|
||||
let db_update_tx = db_update_tx.clone();
|
||||
_parsing_files_tasks.push(cx.background().spawn(async move {
|
||||
let mut retriever = CodeContextRetriever::new();
|
||||
while let Ok(pending_file) = parsing_files_rx.recv().await {
|
||||
Self::parse_file(
|
||||
&fs,
|
||||
pending_file,
|
||||
&mut retriever,
|
||||
&batch_files_tx,
|
||||
&parsing_files_rx,
|
||||
&db_update_tx,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}));
|
||||
}
|
||||
|
||||
log::trace!(
|
||||
"semantic index task initialization took {:?} milliseconds",
|
||||
t0.elapsed().as_millis()
|
||||
);
|
||||
Self {
|
||||
fs,
|
||||
database_url,
|
||||
embedding_provider,
|
||||
language_registry,
|
||||
db_update_tx,
|
||||
parsing_files_tx,
|
||||
_db_update_task,
|
||||
_embed_batch_tasks,
|
||||
_batch_files_task,
|
||||
_parsing_files_tasks,
|
||||
projects: HashMap::new(),
|
||||
}
|
||||
}))
|
||||
}
|
||||
|
||||
fn run_db_operation(db: &VectorDatabase, job: DbOperation) {
|
||||
match job {
|
||||
DbOperation::InsertFile {
|
||||
worktree_id,
|
||||
documents,
|
||||
path,
|
||||
mtime,
|
||||
job_handle,
|
||||
} => {
|
||||
db.insert_file(worktree_id, path, mtime, documents)
|
||||
.log_err();
|
||||
drop(job_handle)
|
||||
}
|
||||
DbOperation::Delete { worktree_id, path } => {
|
||||
db.delete_file(worktree_id, path).log_err();
|
||||
}
|
||||
DbOperation::FindOrCreateWorktree { path, sender } => {
|
||||
let id = db.find_or_create_worktree(&path);
|
||||
sender.send(id).ok();
|
||||
}
|
||||
DbOperation::FileMTimes {
|
||||
worktree_id: worktree_db_id,
|
||||
sender,
|
||||
} => {
|
||||
let file_mtimes = db.get_file_mtimes(worktree_db_id);
|
||||
sender.send(file_mtimes).ok();
|
||||
}
|
||||
DbOperation::WorktreePreviouslyIndexed { path, sender } => {
|
||||
let worktree_indexed = db.worktree_previously_indexed(path.as_ref());
|
||||
sender.send(worktree_indexed).ok();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn compute_embeddings_for_batch(
|
||||
mut embeddings_queue: Vec<(i64, Vec<Document>, PathBuf, SystemTime, JobHandle)>,
|
||||
embedding_provider: &Arc<dyn EmbeddingProvider>,
|
||||
db_update_tx: &channel::Sender<DbOperation>,
|
||||
) {
|
||||
let mut batch_documents = vec![];
|
||||
for (_, documents, _, _, _) in embeddings_queue.iter() {
|
||||
batch_documents.extend(documents.iter().map(|document| document.content.as_str()));
|
||||
}
|
||||
|
||||
if let Ok(embeddings) = embedding_provider.embed_batch(batch_documents).await {
|
||||
log::trace!(
|
||||
"created {} embeddings for {} files",
|
||||
embeddings.len(),
|
||||
embeddings_queue.len(),
|
||||
);
|
||||
|
||||
let mut i = 0;
|
||||
let mut j = 0;
|
||||
|
||||
for embedding in embeddings.iter() {
|
||||
while embeddings_queue[i].1.len() == j {
|
||||
i += 1;
|
||||
j = 0;
|
||||
}
|
||||
|
||||
embeddings_queue[i].1[j].embedding = embedding.to_owned();
|
||||
j += 1;
|
||||
}
|
||||
|
||||
for (worktree_id, documents, path, mtime, job_handle) in embeddings_queue.into_iter() {
|
||||
db_update_tx
|
||||
.send(DbOperation::InsertFile {
|
||||
worktree_id,
|
||||
documents,
|
||||
path,
|
||||
mtime,
|
||||
job_handle,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn enqueue_documents_to_embed(
|
||||
job: EmbeddingJob,
|
||||
queue_len: &mut usize,
|
||||
embeddings_queue: &mut Vec<(i64, Vec<Document>, PathBuf, SystemTime, JobHandle)>,
|
||||
embed_batch_tx: &channel::Sender<Vec<(i64, Vec<Document>, PathBuf, SystemTime, JobHandle)>>,
|
||||
) {
|
||||
let should_flush = match job {
|
||||
EmbeddingJob::Enqueue {
|
||||
documents,
|
||||
worktree_id,
|
||||
path,
|
||||
mtime,
|
||||
job_handle,
|
||||
} => {
|
||||
*queue_len += &documents.len();
|
||||
embeddings_queue.push((worktree_id, documents, path, mtime, job_handle));
|
||||
*queue_len >= EMBEDDINGS_BATCH_SIZE
|
||||
}
|
||||
EmbeddingJob::Flush => true,
|
||||
};
|
||||
|
||||
if should_flush {
|
||||
embed_batch_tx
|
||||
.try_send(mem::take(embeddings_queue))
|
||||
.unwrap();
|
||||
*queue_len = 0;
|
||||
}
|
||||
}
|
||||
|
||||
async fn parse_file(
|
||||
fs: &Arc<dyn Fs>,
|
||||
pending_file: PendingFile,
|
||||
retriever: &mut CodeContextRetriever,
|
||||
batch_files_tx: &channel::Sender<EmbeddingJob>,
|
||||
parsing_files_rx: &channel::Receiver<PendingFile>,
|
||||
db_update_tx: &channel::Sender<DbOperation>,
|
||||
) {
|
||||
if let Some(content) = fs.load(&pending_file.absolute_path).await.log_err() {
|
||||
if let Some(documents) = retriever
|
||||
.parse_file_with_template(
|
||||
&pending_file.relative_path,
|
||||
&content,
|
||||
pending_file.language,
|
||||
)
|
||||
.log_err()
|
||||
{
|
||||
log::trace!(
|
||||
"parsed path {:?}: {} documents",
|
||||
pending_file.relative_path,
|
||||
documents.len()
|
||||
);
|
||||
|
||||
if documents.len() == 0 {
|
||||
db_update_tx
|
||||
.send(DbOperation::InsertFile {
|
||||
worktree_id: pending_file.worktree_db_id,
|
||||
documents,
|
||||
path: pending_file.relative_path,
|
||||
mtime: pending_file.modified_time,
|
||||
job_handle: pending_file.job_handle,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
} else {
|
||||
batch_files_tx
|
||||
.try_send(EmbeddingJob::Enqueue {
|
||||
worktree_id: pending_file.worktree_db_id,
|
||||
path: pending_file.relative_path,
|
||||
mtime: pending_file.modified_time,
|
||||
job_handle: pending_file.job_handle,
|
||||
documents,
|
||||
})
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if parsing_files_rx.len() == 0 {
|
||||
batch_files_tx.try_send(EmbeddingJob::Flush).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
fn find_or_create_worktree(&self, path: PathBuf) -> impl Future<Output = Result<i64>> {
|
||||
let (tx, rx) = oneshot::channel();
|
||||
self.db_update_tx
|
||||
.try_send(DbOperation::FindOrCreateWorktree { path, sender: tx })
|
||||
.unwrap();
|
||||
async move { rx.await? }
|
||||
}
|
||||
|
||||
fn get_file_mtimes(
|
||||
&self,
|
||||
worktree_id: i64,
|
||||
) -> impl Future<Output = Result<HashMap<PathBuf, SystemTime>>> {
|
||||
let (tx, rx) = oneshot::channel();
|
||||
self.db_update_tx
|
||||
.try_send(DbOperation::FileMTimes {
|
||||
worktree_id,
|
||||
sender: tx,
|
||||
})
|
||||
.unwrap();
|
||||
async move { rx.await? }
|
||||
}
|
||||
|
||||
fn worktree_previously_indexed(&self, path: Arc<Path>) -> impl Future<Output = Result<bool>> {
|
||||
let (tx, rx) = oneshot::channel();
|
||||
self.db_update_tx
|
||||
.try_send(DbOperation::WorktreePreviouslyIndexed { path, sender: tx })
|
||||
.unwrap();
|
||||
async move { rx.await? }
|
||||
}
|
||||
|
||||
pub fn project_previously_indexed(
|
||||
&mut self,
|
||||
project: ModelHandle<Project>,
|
||||
cx: &mut ModelContext<Self>,
|
||||
) -> Task<Result<bool>> {
|
||||
let worktree_scans_complete = project
|
||||
.read(cx)
|
||||
.worktrees(cx)
|
||||
.map(|worktree| {
|
||||
let scan_complete = worktree.read(cx).as_local().unwrap().scan_complete();
|
||||
async move {
|
||||
scan_complete.await;
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let worktrees_indexed_previously = project
|
||||
.read(cx)
|
||||
.worktrees(cx)
|
||||
.map(|worktree| self.worktree_previously_indexed(worktree.read(cx).abs_path()))
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
cx.spawn(|_, _cx| async move {
|
||||
futures::future::join_all(worktree_scans_complete).await;
|
||||
|
||||
let worktree_indexed_previously =
|
||||
futures::future::join_all(worktrees_indexed_previously).await;
|
||||
|
||||
Ok(worktree_indexed_previously
|
||||
.iter()
|
||||
.filter(|worktree| worktree.is_ok())
|
||||
.all(|v| v.as_ref().log_err().is_some_and(|v| v.to_owned())))
|
||||
})
|
||||
}
|
||||
|
||||
pub fn index_project(
|
||||
&mut self,
|
||||
project: ModelHandle<Project>,
|
||||
cx: &mut ModelContext<Self>,
|
||||
) -> Task<Result<(usize, watch::Receiver<usize>)>> {
|
||||
let t0 = Instant::now();
|
||||
let worktree_scans_complete = project
|
||||
.read(cx)
|
||||
.worktrees(cx)
|
||||
.map(|worktree| {
|
||||
let scan_complete = worktree.read(cx).as_local().unwrap().scan_complete();
|
||||
async move {
|
||||
scan_complete.await;
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let worktree_db_ids = project
|
||||
.read(cx)
|
||||
.worktrees(cx)
|
||||
.map(|worktree| {
|
||||
self.find_or_create_worktree(worktree.read(cx).abs_path().to_path_buf())
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let language_registry = self.language_registry.clone();
|
||||
let db_update_tx = self.db_update_tx.clone();
|
||||
let parsing_files_tx = self.parsing_files_tx.clone();
|
||||
|
||||
cx.spawn(|this, mut cx| async move {
|
||||
futures::future::join_all(worktree_scans_complete).await;
|
||||
|
||||
let worktree_db_ids = futures::future::join_all(worktree_db_ids).await;
|
||||
|
||||
let worktrees = project.read_with(&cx, |project, cx| {
|
||||
project
|
||||
.worktrees(cx)
|
||||
.map(|worktree| worktree.read(cx).snapshot())
|
||||
.collect::<Vec<_>>()
|
||||
});
|
||||
|
||||
let mut worktree_file_mtimes = HashMap::new();
|
||||
let mut db_ids_by_worktree_id = HashMap::new();
|
||||
for (worktree, db_id) in worktrees.iter().zip(worktree_db_ids) {
|
||||
let db_id = db_id?;
|
||||
db_ids_by_worktree_id.insert(worktree.id(), db_id);
|
||||
worktree_file_mtimes.insert(
|
||||
worktree.id(),
|
||||
this.read_with(&cx, |this, _| this.get_file_mtimes(db_id))
|
||||
.await?,
|
||||
);
|
||||
}
|
||||
|
||||
let (job_count_tx, job_count_rx) = watch::channel_with(0);
|
||||
let job_count_tx = Arc::new(Mutex::new(job_count_tx));
|
||||
this.update(&mut cx, |this, _| {
|
||||
this.projects.insert(
|
||||
project.downgrade(),
|
||||
ProjectState {
|
||||
worktree_db_ids: db_ids_by_worktree_id
|
||||
.iter()
|
||||
.map(|(a, b)| (*a, *b))
|
||||
.collect(),
|
||||
outstanding_job_count_rx: job_count_rx.clone(),
|
||||
_outstanding_job_count_tx: job_count_tx.clone(),
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
cx.background()
|
||||
.spawn(async move {
|
||||
let mut count = 0;
|
||||
for worktree in worktrees.into_iter() {
|
||||
let mut file_mtimes = worktree_file_mtimes.remove(&worktree.id()).unwrap();
|
||||
for file in worktree.files(false, 0) {
|
||||
let absolute_path = worktree.absolutize(&file.path);
|
||||
|
||||
if let Ok(language) = language_registry
|
||||
.language_for_file(&absolute_path, None)
|
||||
.await
|
||||
{
|
||||
if !PARSEABLE_ENTIRE_FILE_TYPES.contains(&language.name().as_ref())
|
||||
&& language
|
||||
.grammar()
|
||||
.and_then(|grammar| grammar.embedding_config.as_ref())
|
||||
.is_none()
|
||||
{
|
||||
continue;
|
||||
}
|
||||
|
||||
let path_buf = file.path.to_path_buf();
|
||||
let stored_mtime = file_mtimes.remove(&file.path.to_path_buf());
|
||||
let already_stored = stored_mtime
|
||||
.map_or(false, |existing_mtime| existing_mtime == file.mtime);
|
||||
|
||||
if !already_stored {
|
||||
count += 1;
|
||||
*job_count_tx.lock().borrow_mut() += 1;
|
||||
let job_handle = JobHandle {
|
||||
tx: Arc::downgrade(&job_count_tx),
|
||||
};
|
||||
parsing_files_tx
|
||||
.try_send(PendingFile {
|
||||
worktree_db_id: db_ids_by_worktree_id[&worktree.id()],
|
||||
relative_path: path_buf,
|
||||
absolute_path,
|
||||
language,
|
||||
job_handle,
|
||||
modified_time: file.mtime,
|
||||
})
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
}
|
||||
for file in file_mtimes.keys() {
|
||||
db_update_tx
|
||||
.try_send(DbOperation::Delete {
|
||||
worktree_id: db_ids_by_worktree_id[&worktree.id()],
|
||||
path: file.to_owned(),
|
||||
})
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
log::trace!(
|
||||
"walking worktree took {:?} milliseconds",
|
||||
t0.elapsed().as_millis()
|
||||
);
|
||||
anyhow::Ok((count, job_count_rx))
|
||||
})
|
||||
.await
|
||||
})
|
||||
}
|
||||
|
||||
pub fn outstanding_job_count_rx(
|
||||
&self,
|
||||
project: &ModelHandle<Project>,
|
||||
) -> Option<watch::Receiver<usize>> {
|
||||
Some(
|
||||
self.projects
|
||||
.get(&project.downgrade())?
|
||||
.outstanding_job_count_rx
|
||||
.clone(),
|
||||
)
|
||||
}
|
||||
|
||||
pub fn search_project(
|
||||
&mut self,
|
||||
project: ModelHandle<Project>,
|
||||
phrase: String,
|
||||
limit: usize,
|
||||
include_globs: Vec<GlobMatcher>,
|
||||
exclude_globs: Vec<GlobMatcher>,
|
||||
cx: &mut ModelContext<Self>,
|
||||
) -> Task<Result<Vec<SearchResult>>> {
|
||||
let project_state = if let Some(state) = self.projects.get(&project.downgrade()) {
|
||||
state
|
||||
} else {
|
||||
return Task::ready(Err(anyhow!("project not added")));
|
||||
};
|
||||
|
||||
let worktree_db_ids = project
|
||||
.read(cx)
|
||||
.worktrees(cx)
|
||||
.filter_map(|worktree| {
|
||||
let worktree_id = worktree.read(cx).id();
|
||||
project_state.db_id_for_worktree_id(worktree_id)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let embedding_provider = self.embedding_provider.clone();
|
||||
let database_url = self.database_url.clone();
|
||||
let fs = self.fs.clone();
|
||||
cx.spawn(|this, mut cx| async move {
|
||||
let documents = cx
|
||||
.background()
|
||||
.spawn(async move {
|
||||
let database = VectorDatabase::new(fs, database_url).await?;
|
||||
|
||||
let phrase_embedding = embedding_provider
|
||||
.embed_batch(vec![&phrase])
|
||||
.await?
|
||||
.into_iter()
|
||||
.next()
|
||||
.unwrap();
|
||||
|
||||
database.top_k_search(
|
||||
&worktree_db_ids,
|
||||
&phrase_embedding,
|
||||
limit,
|
||||
include_globs,
|
||||
exclude_globs,
|
||||
)
|
||||
})
|
||||
.await?;
|
||||
|
||||
let mut tasks = Vec::new();
|
||||
let mut ranges = Vec::new();
|
||||
let weak_project = project.downgrade();
|
||||
project.update(&mut cx, |project, cx| {
|
||||
for (worktree_db_id, file_path, byte_range) in documents {
|
||||
let project_state =
|
||||
if let Some(state) = this.read(cx).projects.get(&weak_project) {
|
||||
state
|
||||
} else {
|
||||
return Err(anyhow!("project not added"));
|
||||
};
|
||||
if let Some(worktree_id) = project_state.worktree_id_for_db_id(worktree_db_id) {
|
||||
tasks.push(project.open_buffer((worktree_id, file_path), cx));
|
||||
ranges.push(byte_range);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
})?;
|
||||
|
||||
let buffers = futures::future::join_all(tasks).await;
|
||||
|
||||
Ok(buffers
|
||||
.into_iter()
|
||||
.zip(ranges)
|
||||
.filter_map(|(buffer, range)| {
|
||||
let buffer = buffer.log_err()?;
|
||||
let range = buffer.read_with(&cx, |buffer, _| {
|
||||
buffer.anchor_before(range.start)..buffer.anchor_after(range.end)
|
||||
});
|
||||
Some(SearchResult { buffer, range })
|
||||
})
|
||||
.collect::<Vec<_>>())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl Entity for SemanticIndex {
|
||||
type Event = ();
|
||||
}
|
||||
|
||||
impl Drop for JobHandle {
|
||||
fn drop(&mut self) {
|
||||
if let Some(tx) = self.tx.upgrade() {
|
||||
let mut tx = tx.lock();
|
||||
*tx.borrow_mut() -= 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
+5
-5
@@ -4,21 +4,21 @@ use serde::{Deserialize, Serialize};
|
||||
use settings::Setting;
|
||||
|
||||
#[derive(Deserialize, Debug)]
|
||||
pub struct VectorStoreSettings {
|
||||
pub struct SemanticIndexSettings {
|
||||
pub enabled: bool,
|
||||
pub reindexing_delay_seconds: usize,
|
||||
}
|
||||
|
||||
#[derive(Clone, Default, Serialize, Deserialize, JsonSchema, Debug)]
|
||||
pub struct VectorStoreSettingsContent {
|
||||
pub struct SemanticIndexSettingsContent {
|
||||
pub enabled: Option<bool>,
|
||||
pub reindexing_delay_seconds: Option<usize>,
|
||||
}
|
||||
|
||||
impl Setting for VectorStoreSettings {
|
||||
const KEY: Option<&'static str> = Some("vector_store");
|
||||
impl Setting for SemanticIndexSettings {
|
||||
const KEY: Option<&'static str> = Some("semantic_index");
|
||||
|
||||
type FileContent = VectorStoreSettingsContent;
|
||||
type FileContent = SemanticIndexSettingsContent;
|
||||
|
||||
fn load(
|
||||
default_value: &Self::FileContent,
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,172 +0,0 @@
|
||||
use crate::{SearchResult, VectorStore};
|
||||
use editor::{scroll::autoscroll::Autoscroll, Editor};
|
||||
use gpui::{
|
||||
actions, elements::*, AnyElement, AppContext, ModelHandle, MouseState, Task, ViewContext,
|
||||
WeakViewHandle,
|
||||
};
|
||||
use picker::{Picker, PickerDelegate, PickerEvent};
|
||||
use project::{Project, ProjectPath};
|
||||
use std::{collections::HashMap, sync::Arc, time::Duration};
|
||||
use util::ResultExt;
|
||||
use workspace::Workspace;
|
||||
|
||||
const MIN_QUERY_LEN: usize = 5;
|
||||
const EMBEDDING_DEBOUNCE_INTERVAL: Duration = Duration::from_millis(500);
|
||||
|
||||
actions!(semantic_search, [Toggle]);
|
||||
|
||||
pub type SemanticSearch = Picker<SemanticSearchDelegate>;
|
||||
|
||||
pub struct SemanticSearchDelegate {
|
||||
workspace: WeakViewHandle<Workspace>,
|
||||
project: ModelHandle<Project>,
|
||||
vector_store: ModelHandle<VectorStore>,
|
||||
selected_match_index: usize,
|
||||
matches: Vec<SearchResult>,
|
||||
history: HashMap<String, Vec<SearchResult>>,
|
||||
}
|
||||
|
||||
impl SemanticSearchDelegate {
|
||||
// This is currently searching on every keystroke,
|
||||
// This is wildly overkill, and has the potential to get expensive
|
||||
// We will need to update this to throttle searching
|
||||
pub fn new(
|
||||
workspace: WeakViewHandle<Workspace>,
|
||||
project: ModelHandle<Project>,
|
||||
vector_store: ModelHandle<VectorStore>,
|
||||
) -> Self {
|
||||
Self {
|
||||
workspace,
|
||||
project,
|
||||
vector_store,
|
||||
selected_match_index: 0,
|
||||
matches: vec![],
|
||||
history: HashMap::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl PickerDelegate for SemanticSearchDelegate {
|
||||
fn placeholder_text(&self) -> Arc<str> {
|
||||
"Search repository in natural language...".into()
|
||||
}
|
||||
|
||||
fn confirm(&mut self, _: bool, cx: &mut ViewContext<SemanticSearch>) {
|
||||
if let Some(search_result) = self.matches.get(self.selected_match_index) {
|
||||
// Open Buffer
|
||||
let search_result = search_result.clone();
|
||||
let buffer = self.project.update(cx, |project, cx| {
|
||||
project.open_buffer(
|
||||
ProjectPath {
|
||||
worktree_id: search_result.worktree_id,
|
||||
path: search_result.file_path.clone().into(),
|
||||
},
|
||||
cx,
|
||||
)
|
||||
});
|
||||
|
||||
let workspace = self.workspace.clone();
|
||||
let position = search_result.clone().offset;
|
||||
cx.spawn(|_, mut cx| async move {
|
||||
let buffer = buffer.await?;
|
||||
workspace.update(&mut cx, |workspace, cx| {
|
||||
let editor = workspace.open_project_item::<Editor>(buffer, cx);
|
||||
editor.update(cx, |editor, cx| {
|
||||
editor.change_selections(Some(Autoscroll::center()), cx, |s| {
|
||||
s.select_ranges([position..position])
|
||||
});
|
||||
});
|
||||
})?;
|
||||
Ok::<_, anyhow::Error>(())
|
||||
})
|
||||
.detach_and_log_err(cx);
|
||||
cx.emit(PickerEvent::Dismiss);
|
||||
}
|
||||
}
|
||||
|
||||
fn dismissed(&mut self, _cx: &mut ViewContext<SemanticSearch>) {}
|
||||
|
||||
fn match_count(&self) -> usize {
|
||||
self.matches.len()
|
||||
}
|
||||
|
||||
fn selected_index(&self) -> usize {
|
||||
self.selected_match_index
|
||||
}
|
||||
|
||||
fn set_selected_index(&mut self, ix: usize, _cx: &mut ViewContext<SemanticSearch>) {
|
||||
self.selected_match_index = ix;
|
||||
}
|
||||
|
||||
fn update_matches(&mut self, query: String, cx: &mut ViewContext<SemanticSearch>) -> Task<()> {
|
||||
log::info!("Searching for {:?}...", query);
|
||||
if query.len() < MIN_QUERY_LEN {
|
||||
log::info!("Query below minimum length");
|
||||
return Task::ready(());
|
||||
}
|
||||
|
||||
let vector_store = self.vector_store.clone();
|
||||
let project = self.project.clone();
|
||||
cx.spawn(|this, mut cx| async move {
|
||||
cx.background().timer(EMBEDDING_DEBOUNCE_INTERVAL).await;
|
||||
|
||||
let retrieved_cached = this.update(&mut cx, |this, _| {
|
||||
let delegate = this.delegate_mut();
|
||||
if delegate.history.contains_key(&query) {
|
||||
let historic_results = delegate.history.get(&query).unwrap().to_owned();
|
||||
delegate.matches = historic_results.clone();
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
});
|
||||
|
||||
if let Some(retrieved) = retrieved_cached.log_err() {
|
||||
if !retrieved {
|
||||
let task = vector_store.update(&mut cx, |store, cx| {
|
||||
store.search(project.clone(), query.to_string(), 10, cx)
|
||||
});
|
||||
|
||||
if let Some(results) = task.await.log_err() {
|
||||
log::info!("Not queried previously, searching...");
|
||||
this.update(&mut cx, |this, _| {
|
||||
let delegate = this.delegate_mut();
|
||||
delegate.matches = results.clone();
|
||||
delegate.history.insert(query, results);
|
||||
})
|
||||
.ok();
|
||||
}
|
||||
} else {
|
||||
log::info!("Already queried, retrieved directly from cached history");
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn render_match(
|
||||
&self,
|
||||
ix: usize,
|
||||
mouse_state: &mut MouseState,
|
||||
selected: bool,
|
||||
cx: &AppContext,
|
||||
) -> AnyElement<Picker<Self>> {
|
||||
let theme = theme::current(cx);
|
||||
let style = &theme.picker.item;
|
||||
let current_style = style.in_state(selected).style_for(mouse_state);
|
||||
|
||||
let search_result = &self.matches[ix];
|
||||
|
||||
let path = search_result.file_path.to_string_lossy();
|
||||
let name = search_result.name.clone();
|
||||
|
||||
Flex::column()
|
||||
.with_child(Text::new(name, current_style.label.text.clone()).with_soft_wrap(false))
|
||||
.with_child(Label::new(
|
||||
path.to_string(),
|
||||
style.inactive_state().default.label.clone(),
|
||||
))
|
||||
.contained()
|
||||
.with_style(current_style.container)
|
||||
.into_any()
|
||||
}
|
||||
}
|
||||
@@ -1,115 +0,0 @@
|
||||
use std::{path::PathBuf, sync::Arc, time::SystemTime};
|
||||
|
||||
use anyhow::{anyhow, Ok, Result};
|
||||
use project::Fs;
|
||||
use tree_sitter::{Parser, QueryCursor};
|
||||
|
||||
use crate::PendingFile;
|
||||
|
||||
#[derive(Debug, PartialEq, Clone)]
|
||||
pub struct Document {
|
||||
pub offset: usize,
|
||||
pub name: String,
|
||||
pub embedding: Vec<f32>,
|
||||
}
|
||||
|
||||
#[derive(Debug, PartialEq, Clone)]
|
||||
pub struct ParsedFile {
|
||||
pub path: PathBuf,
|
||||
pub mtime: SystemTime,
|
||||
pub documents: Vec<Document>,
|
||||
}
|
||||
|
||||
const CODE_CONTEXT_TEMPLATE: &str =
|
||||
"The below code snippet is from file '<path>'\n\n```<language>\n<item>\n```";
|
||||
|
||||
pub struct CodeContextRetriever {
|
||||
pub parser: Parser,
|
||||
pub cursor: QueryCursor,
|
||||
pub fs: Arc<dyn Fs>,
|
||||
}
|
||||
|
||||
impl CodeContextRetriever {
|
||||
pub async fn parse_file(
|
||||
&mut self,
|
||||
pending_file: PendingFile,
|
||||
) -> Result<(ParsedFile, Vec<String>)> {
|
||||
let grammar = pending_file
|
||||
.language
|
||||
.grammar()
|
||||
.ok_or_else(|| anyhow!("no grammar for language"))?;
|
||||
let embedding_config = grammar
|
||||
.embedding_config
|
||||
.as_ref()
|
||||
.ok_or_else(|| anyhow!("no embedding queries"))?;
|
||||
|
||||
let content = self.fs.load(&pending_file.absolute_path).await?;
|
||||
|
||||
self.parser.set_language(grammar.ts_language).unwrap();
|
||||
|
||||
let tree = self
|
||||
.parser
|
||||
.parse(&content, None)
|
||||
.ok_or_else(|| anyhow!("parsing failed"))?;
|
||||
|
||||
let mut documents = Vec::new();
|
||||
let mut context_spans = Vec::new();
|
||||
|
||||
// Iterate through query matches
|
||||
for mat in self.cursor.matches(
|
||||
&embedding_config.query,
|
||||
tree.root_node(),
|
||||
content.as_bytes(),
|
||||
) {
|
||||
// log::info!("-----MATCH-----");
|
||||
|
||||
let mut name = Vec::new();
|
||||
let mut item: Option<&str> = None;
|
||||
let mut offset: Option<usize> = None;
|
||||
for capture in mat.captures {
|
||||
if capture.index == embedding_config.item_capture_ix {
|
||||
offset = Some(capture.node.byte_range().start);
|
||||
item = content.get(capture.node.byte_range());
|
||||
} else if capture.index == embedding_config.name_capture_ix {
|
||||
if let Some(name_content) = content.get(capture.node.byte_range()) {
|
||||
name.push(name_content);
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(context_capture_ix) = embedding_config.context_capture_ix {
|
||||
if capture.index == context_capture_ix {
|
||||
if let Some(context) = content.get(capture.node.byte_range()) {
|
||||
name.push(context);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if item.is_some() && offset.is_some() && name.len() > 0 {
|
||||
let context_span = CODE_CONTEXT_TEMPLATE
|
||||
.replace("<path>", pending_file.relative_path.to_str().unwrap())
|
||||
.replace("<language>", &pending_file.language.name().to_lowercase())
|
||||
.replace("<item>", item.unwrap());
|
||||
|
||||
// log::info!("Name: {:?}", name);
|
||||
// log::info!("Span: {:?}", util::truncate(&context_span, 100));
|
||||
|
||||
context_spans.push(context_span);
|
||||
documents.push(Document {
|
||||
name: name.join(" "),
|
||||
offset: offset.unwrap(),
|
||||
embedding: Vec::new(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
return Ok((
|
||||
ParsedFile {
|
||||
path: pending_file.relative_path,
|
||||
mtime: pending_file.modified_time,
|
||||
documents,
|
||||
},
|
||||
context_spans,
|
||||
));
|
||||
}
|
||||
}
|
||||
@@ -1,770 +0,0 @@
|
||||
mod db;
|
||||
mod embedding;
|
||||
mod modal;
|
||||
mod parsing;
|
||||
mod vector_store_settings;
|
||||
|
||||
#[cfg(test)]
|
||||
mod vector_store_tests;
|
||||
|
||||
use crate::vector_store_settings::VectorStoreSettings;
|
||||
use anyhow::{anyhow, Result};
|
||||
use db::VectorDatabase;
|
||||
use embedding::{EmbeddingProvider, OpenAIEmbeddings};
|
||||
use futures::{channel::oneshot, Future};
|
||||
use gpui::{
|
||||
AppContext, AsyncAppContext, Entity, ModelContext, ModelHandle, Task, ViewContext,
|
||||
WeakModelHandle,
|
||||
};
|
||||
use language::{Language, LanguageRegistry};
|
||||
use modal::{SemanticSearch, SemanticSearchDelegate, Toggle};
|
||||
use parsing::{CodeContextRetriever, ParsedFile};
|
||||
use project::{Fs, PathChange, Project, ProjectEntryId, WorktreeId};
|
||||
use smol::channel;
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
path::{Path, PathBuf},
|
||||
sync::Arc,
|
||||
time::{Duration, Instant, SystemTime},
|
||||
};
|
||||
use tree_sitter::{Parser, QueryCursor};
|
||||
use util::{
|
||||
channel::{ReleaseChannel, RELEASE_CHANNEL, RELEASE_CHANNEL_NAME},
|
||||
http::HttpClient,
|
||||
paths::EMBEDDINGS_DIR,
|
||||
ResultExt,
|
||||
};
|
||||
use workspace::{Workspace, WorkspaceCreated};
|
||||
|
||||
const VECTOR_STORE_VERSION: usize = 0;
|
||||
const EMBEDDINGS_BATCH_SIZE: usize = 150;
|
||||
|
||||
pub fn init(
|
||||
fs: Arc<dyn Fs>,
|
||||
http_client: Arc<dyn HttpClient>,
|
||||
language_registry: Arc<LanguageRegistry>,
|
||||
cx: &mut AppContext,
|
||||
) {
|
||||
settings::register::<VectorStoreSettings>(cx);
|
||||
|
||||
let db_file_path = EMBEDDINGS_DIR
|
||||
.join(Path::new(RELEASE_CHANNEL_NAME.as_str()))
|
||||
.join("embeddings_db");
|
||||
|
||||
SemanticSearch::init(cx);
|
||||
cx.add_action(
|
||||
|workspace: &mut Workspace, _: &Toggle, cx: &mut ViewContext<Workspace>| {
|
||||
if cx.has_global::<ModelHandle<VectorStore>>() {
|
||||
let vector_store = cx.global::<ModelHandle<VectorStore>>().clone();
|
||||
workspace.toggle_modal(cx, |workspace, cx| {
|
||||
let project = workspace.project().clone();
|
||||
let workspace = cx.weak_handle();
|
||||
cx.add_view(|cx| {
|
||||
SemanticSearch::new(
|
||||
SemanticSearchDelegate::new(workspace, project, vector_store),
|
||||
cx,
|
||||
)
|
||||
})
|
||||
});
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
if *RELEASE_CHANNEL == ReleaseChannel::Stable
|
||||
|| !settings::get::<VectorStoreSettings>(cx).enabled
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
cx.spawn(move |mut cx| async move {
|
||||
let vector_store = VectorStore::new(
|
||||
fs,
|
||||
db_file_path,
|
||||
// Arc::new(embedding::DummyEmbeddings {}),
|
||||
Arc::new(OpenAIEmbeddings {
|
||||
client: http_client,
|
||||
executor: cx.background(),
|
||||
}),
|
||||
language_registry,
|
||||
cx.clone(),
|
||||
)
|
||||
.await?;
|
||||
|
||||
cx.update(|cx| {
|
||||
cx.set_global(vector_store.clone());
|
||||
cx.subscribe_global::<WorkspaceCreated, _>({
|
||||
let vector_store = vector_store.clone();
|
||||
move |event, cx| {
|
||||
let workspace = &event.0;
|
||||
if let Some(workspace) = workspace.upgrade(cx) {
|
||||
let project = workspace.read(cx).project().clone();
|
||||
if project.read(cx).is_local() {
|
||||
vector_store.update(cx, |store, cx| {
|
||||
store.add_project(project, cx).detach();
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
.detach();
|
||||
});
|
||||
|
||||
anyhow::Ok(())
|
||||
})
|
||||
.detach();
|
||||
}
|
||||
|
||||
pub struct VectorStore {
|
||||
fs: Arc<dyn Fs>,
|
||||
database_url: Arc<PathBuf>,
|
||||
embedding_provider: Arc<dyn EmbeddingProvider>,
|
||||
language_registry: Arc<LanguageRegistry>,
|
||||
db_update_tx: channel::Sender<DbOperation>,
|
||||
parsing_files_tx: channel::Sender<PendingFile>,
|
||||
_db_update_task: Task<()>,
|
||||
_embed_batch_task: Task<()>,
|
||||
_batch_files_task: Task<()>,
|
||||
_parsing_files_tasks: Vec<Task<()>>,
|
||||
projects: HashMap<WeakModelHandle<Project>, ProjectState>,
|
||||
}
|
||||
|
||||
struct ProjectState {
|
||||
worktree_db_ids: Vec<(WorktreeId, i64)>,
|
||||
pending_files: HashMap<PathBuf, (PendingFile, SystemTime)>,
|
||||
_subscription: gpui::Subscription,
|
||||
}
|
||||
|
||||
impl ProjectState {
|
||||
fn db_id_for_worktree_id(&self, id: WorktreeId) -> Option<i64> {
|
||||
self.worktree_db_ids
|
||||
.iter()
|
||||
.find_map(|(worktree_id, db_id)| {
|
||||
if *worktree_id == id {
|
||||
Some(*db_id)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn worktree_id_for_db_id(&self, id: i64) -> Option<WorktreeId> {
|
||||
self.worktree_db_ids
|
||||
.iter()
|
||||
.find_map(|(worktree_id, db_id)| {
|
||||
if *db_id == id {
|
||||
Some(*worktree_id)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn update_pending_files(&mut self, pending_file: PendingFile, indexing_time: SystemTime) {
|
||||
// If Pending File Already Exists, Replace it with the new one
|
||||
// but keep the old indexing time
|
||||
if let Some(old_file) = self
|
||||
.pending_files
|
||||
.remove(&pending_file.relative_path.clone())
|
||||
{
|
||||
self.pending_files.insert(
|
||||
pending_file.relative_path.clone(),
|
||||
(pending_file, old_file.1),
|
||||
);
|
||||
} else {
|
||||
self.pending_files.insert(
|
||||
pending_file.relative_path.clone(),
|
||||
(pending_file, indexing_time),
|
||||
);
|
||||
};
|
||||
}
|
||||
|
||||
fn get_outstanding_files(&mut self) -> Vec<PendingFile> {
|
||||
let mut outstanding_files = vec![];
|
||||
let mut remove_keys = vec![];
|
||||
for key in self.pending_files.keys().into_iter() {
|
||||
if let Some(pending_details) = self.pending_files.get(key) {
|
||||
let (pending_file, index_time) = pending_details;
|
||||
if index_time <= &SystemTime::now() {
|
||||
outstanding_files.push(pending_file.clone());
|
||||
remove_keys.push(key.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for key in remove_keys.iter() {
|
||||
self.pending_files.remove(key);
|
||||
}
|
||||
|
||||
return outstanding_files;
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct PendingFile {
|
||||
worktree_db_id: i64,
|
||||
relative_path: PathBuf,
|
||||
absolute_path: PathBuf,
|
||||
language: Arc<Language>,
|
||||
modified_time: SystemTime,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SearchResult {
|
||||
pub worktree_id: WorktreeId,
|
||||
pub name: String,
|
||||
pub offset: usize,
|
||||
pub file_path: PathBuf,
|
||||
}
|
||||
|
||||
enum DbOperation {
|
||||
InsertFile {
|
||||
worktree_id: i64,
|
||||
indexed_file: ParsedFile,
|
||||
},
|
||||
Delete {
|
||||
worktree_id: i64,
|
||||
path: PathBuf,
|
||||
},
|
||||
FindOrCreateWorktree {
|
||||
path: PathBuf,
|
||||
sender: oneshot::Sender<Result<i64>>,
|
||||
},
|
||||
FileMTimes {
|
||||
worktree_id: i64,
|
||||
sender: oneshot::Sender<Result<HashMap<PathBuf, SystemTime>>>,
|
||||
},
|
||||
}
|
||||
|
||||
enum EmbeddingJob {
|
||||
Enqueue {
|
||||
worktree_id: i64,
|
||||
parsed_file: ParsedFile,
|
||||
document_spans: Vec<String>,
|
||||
},
|
||||
Flush,
|
||||
}
|
||||
|
||||
impl VectorStore {
|
||||
async fn new(
|
||||
fs: Arc<dyn Fs>,
|
||||
database_url: PathBuf,
|
||||
embedding_provider: Arc<dyn EmbeddingProvider>,
|
||||
language_registry: Arc<LanguageRegistry>,
|
||||
mut cx: AsyncAppContext,
|
||||
) -> Result<ModelHandle<Self>> {
|
||||
let database_url = Arc::new(database_url);
|
||||
|
||||
let db = cx
|
||||
.background()
|
||||
.spawn({
|
||||
let fs = fs.clone();
|
||||
let database_url = database_url.clone();
|
||||
async move {
|
||||
if let Some(db_directory) = database_url.parent() {
|
||||
fs.create_dir(db_directory).await.log_err();
|
||||
}
|
||||
|
||||
let db = VectorDatabase::new(database_url.to_string_lossy().to_string())?;
|
||||
anyhow::Ok(db)
|
||||
}
|
||||
})
|
||||
.await?;
|
||||
|
||||
Ok(cx.add_model(|cx| {
|
||||
// paths_tx -> embeddings_tx -> db_update_tx
|
||||
|
||||
//db_update_tx/rx: Updating Database
|
||||
let (db_update_tx, db_update_rx) = channel::unbounded();
|
||||
let _db_update_task = cx.background().spawn(async move {
|
||||
while let Ok(job) = db_update_rx.recv().await {
|
||||
match job {
|
||||
DbOperation::InsertFile {
|
||||
worktree_id,
|
||||
indexed_file,
|
||||
} => {
|
||||
db.insert_file(worktree_id, indexed_file).log_err();
|
||||
}
|
||||
DbOperation::Delete { worktree_id, path } => {
|
||||
db.delete_file(worktree_id, path).log_err();
|
||||
}
|
||||
DbOperation::FindOrCreateWorktree { path, sender } => {
|
||||
let id = db.find_or_create_worktree(&path);
|
||||
sender.send(id).ok();
|
||||
}
|
||||
DbOperation::FileMTimes {
|
||||
worktree_id: worktree_db_id,
|
||||
sender,
|
||||
} => {
|
||||
let file_mtimes = db.get_file_mtimes(worktree_db_id);
|
||||
sender.send(file_mtimes).ok();
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// embed_tx/rx: Embed Batch and Send to Database
|
||||
let (embed_batch_tx, embed_batch_rx) =
|
||||
channel::unbounded::<Vec<(i64, ParsedFile, Vec<String>)>>();
|
||||
let _embed_batch_task = cx.background().spawn({
|
||||
let db_update_tx = db_update_tx.clone();
|
||||
let embedding_provider = embedding_provider.clone();
|
||||
async move {
|
||||
while let Ok(mut embeddings_queue) = embed_batch_rx.recv().await {
|
||||
// Construct Batch
|
||||
let mut document_spans = vec![];
|
||||
for (_, _, document_span) in embeddings_queue.iter() {
|
||||
document_spans.extend(document_span.iter().map(|s| s.as_str()));
|
||||
}
|
||||
|
||||
if let Ok(embeddings) = embedding_provider.embed_batch(document_spans).await
|
||||
{
|
||||
let mut i = 0;
|
||||
let mut j = 0;
|
||||
|
||||
for embedding in embeddings.iter() {
|
||||
while embeddings_queue[i].1.documents.len() == j {
|
||||
i += 1;
|
||||
j = 0;
|
||||
}
|
||||
|
||||
embeddings_queue[i].1.documents[j].embedding = embedding.to_owned();
|
||||
j += 1;
|
||||
}
|
||||
|
||||
for (worktree_id, indexed_file, _) in embeddings_queue.into_iter() {
|
||||
for document in indexed_file.documents.iter() {
|
||||
// TODO: Update this so it doesn't panic
|
||||
assert!(
|
||||
document.embedding.len() > 0,
|
||||
"Document Embedding Not Complete"
|
||||
);
|
||||
}
|
||||
|
||||
db_update_tx
|
||||
.send(DbOperation::InsertFile {
|
||||
worktree_id,
|
||||
indexed_file,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// batch_tx/rx: Batch Files to Send for Embeddings
|
||||
let (batch_files_tx, batch_files_rx) = channel::unbounded::<EmbeddingJob>();
|
||||
let _batch_files_task = cx.background().spawn(async move {
|
||||
let mut queue_len = 0;
|
||||
let mut embeddings_queue = vec![];
|
||||
|
||||
while let Ok(job) = batch_files_rx.recv().await {
|
||||
let should_flush = match job {
|
||||
EmbeddingJob::Enqueue {
|
||||
document_spans,
|
||||
worktree_id,
|
||||
parsed_file,
|
||||
} => {
|
||||
queue_len += &document_spans.len();
|
||||
embeddings_queue.push((worktree_id, parsed_file, document_spans));
|
||||
queue_len >= EMBEDDINGS_BATCH_SIZE
|
||||
}
|
||||
EmbeddingJob::Flush => true,
|
||||
};
|
||||
|
||||
if should_flush {
|
||||
embed_batch_tx.try_send(embeddings_queue).unwrap();
|
||||
embeddings_queue = vec![];
|
||||
queue_len = 0;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// parsing_files_tx/rx: Parsing Files to Embeddable Documents
|
||||
let (parsing_files_tx, parsing_files_rx) = channel::unbounded::<PendingFile>();
|
||||
|
||||
let mut _parsing_files_tasks = Vec::new();
|
||||
// for _ in 0..cx.background().num_cpus() {
|
||||
for _ in 0..1 {
|
||||
let fs = fs.clone();
|
||||
let parsing_files_rx = parsing_files_rx.clone();
|
||||
let batch_files_tx = batch_files_tx.clone();
|
||||
_parsing_files_tasks.push(cx.background().spawn(async move {
|
||||
let parser = Parser::new();
|
||||
let cursor = QueryCursor::new();
|
||||
let mut retriever = CodeContextRetriever { parser, cursor, fs };
|
||||
while let Ok(pending_file) = parsing_files_rx.recv().await {
|
||||
if let Some((indexed_file, document_spans)) =
|
||||
retriever.parse_file(pending_file.clone()).await.log_err()
|
||||
{
|
||||
batch_files_tx
|
||||
.try_send(EmbeddingJob::Enqueue {
|
||||
worktree_id: pending_file.worktree_db_id,
|
||||
parsed_file: indexed_file,
|
||||
document_spans,
|
||||
})
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
if parsing_files_rx.len() == 0 {
|
||||
batch_files_tx.try_send(EmbeddingJob::Flush).unwrap();
|
||||
}
|
||||
}
|
||||
}));
|
||||
}
|
||||
|
||||
Self {
|
||||
fs,
|
||||
database_url,
|
||||
embedding_provider,
|
||||
language_registry,
|
||||
db_update_tx,
|
||||
parsing_files_tx,
|
||||
_db_update_task,
|
||||
_embed_batch_task,
|
||||
_batch_files_task,
|
||||
_parsing_files_tasks,
|
||||
projects: HashMap::new(),
|
||||
}
|
||||
}))
|
||||
}
|
||||
|
||||
fn find_or_create_worktree(&self, path: PathBuf) -> impl Future<Output = Result<i64>> {
|
||||
let (tx, rx) = oneshot::channel();
|
||||
self.db_update_tx
|
||||
.try_send(DbOperation::FindOrCreateWorktree { path, sender: tx })
|
||||
.unwrap();
|
||||
async move { rx.await? }
|
||||
}
|
||||
|
||||
fn get_file_mtimes(
|
||||
&self,
|
||||
worktree_id: i64,
|
||||
) -> impl Future<Output = Result<HashMap<PathBuf, SystemTime>>> {
|
||||
let (tx, rx) = oneshot::channel();
|
||||
self.db_update_tx
|
||||
.try_send(DbOperation::FileMTimes {
|
||||
worktree_id,
|
||||
sender: tx,
|
||||
})
|
||||
.unwrap();
|
||||
async move { rx.await? }
|
||||
}
|
||||
|
||||
fn add_project(
|
||||
&mut self,
|
||||
project: ModelHandle<Project>,
|
||||
cx: &mut ModelContext<Self>,
|
||||
) -> Task<Result<()>> {
|
||||
let worktree_scans_complete = project
|
||||
.read(cx)
|
||||
.worktrees(cx)
|
||||
.map(|worktree| {
|
||||
let scan_complete = worktree.read(cx).as_local().unwrap().scan_complete();
|
||||
async move {
|
||||
scan_complete.await;
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let worktree_db_ids = project
|
||||
.read(cx)
|
||||
.worktrees(cx)
|
||||
.map(|worktree| {
|
||||
self.find_or_create_worktree(worktree.read(cx).abs_path().to_path_buf())
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let fs = self.fs.clone();
|
||||
let language_registry = self.language_registry.clone();
|
||||
let database_url = self.database_url.clone();
|
||||
let db_update_tx = self.db_update_tx.clone();
|
||||
let parsing_files_tx = self.parsing_files_tx.clone();
|
||||
|
||||
cx.spawn(|this, mut cx| async move {
|
||||
futures::future::join_all(worktree_scans_complete).await;
|
||||
|
||||
let worktree_db_ids = futures::future::join_all(worktree_db_ids).await;
|
||||
|
||||
if let Some(db_directory) = database_url.parent() {
|
||||
fs.create_dir(db_directory).await.log_err();
|
||||
}
|
||||
|
||||
let worktrees = project.read_with(&cx, |project, cx| {
|
||||
project
|
||||
.worktrees(cx)
|
||||
.map(|worktree| worktree.read(cx).snapshot())
|
||||
.collect::<Vec<_>>()
|
||||
});
|
||||
|
||||
let mut worktree_file_times = HashMap::new();
|
||||
let mut db_ids_by_worktree_id = HashMap::new();
|
||||
for (worktree, db_id) in worktrees.iter().zip(worktree_db_ids) {
|
||||
let db_id = db_id?;
|
||||
db_ids_by_worktree_id.insert(worktree.id(), db_id);
|
||||
worktree_file_times.insert(
|
||||
worktree.id(),
|
||||
this.read_with(&cx, |this, _| this.get_file_mtimes(db_id))
|
||||
.await?,
|
||||
);
|
||||
}
|
||||
|
||||
cx.background()
|
||||
.spawn({
|
||||
let db_ids_by_worktree_id = db_ids_by_worktree_id.clone();
|
||||
let db_update_tx = db_update_tx.clone();
|
||||
let language_registry = language_registry.clone();
|
||||
let parsing_files_tx = parsing_files_tx.clone();
|
||||
async move {
|
||||
let t0 = Instant::now();
|
||||
for worktree in worktrees.into_iter() {
|
||||
let mut file_mtimes =
|
||||
worktree_file_times.remove(&worktree.id()).unwrap();
|
||||
for file in worktree.files(false, 0) {
|
||||
let absolute_path = worktree.absolutize(&file.path);
|
||||
|
||||
if let Ok(language) = language_registry
|
||||
.language_for_file(&absolute_path, None)
|
||||
.await
|
||||
{
|
||||
if language
|
||||
.grammar()
|
||||
.and_then(|grammar| grammar.embedding_config.as_ref())
|
||||
.is_none()
|
||||
{
|
||||
continue;
|
||||
}
|
||||
|
||||
let path_buf = file.path.to_path_buf();
|
||||
let stored_mtime = file_mtimes.remove(&file.path.to_path_buf());
|
||||
let already_stored = stored_mtime
|
||||
.map_or(false, |existing_mtime| {
|
||||
existing_mtime == file.mtime
|
||||
});
|
||||
|
||||
if !already_stored {
|
||||
parsing_files_tx
|
||||
.try_send(PendingFile {
|
||||
worktree_db_id: db_ids_by_worktree_id
|
||||
[&worktree.id()],
|
||||
relative_path: path_buf,
|
||||
absolute_path,
|
||||
language,
|
||||
modified_time: file.mtime,
|
||||
})
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
}
|
||||
for file in file_mtimes.keys() {
|
||||
db_update_tx
|
||||
.try_send(DbOperation::Delete {
|
||||
worktree_id: db_ids_by_worktree_id[&worktree.id()],
|
||||
path: file.to_owned(),
|
||||
})
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
log::info!(
|
||||
"Parsing Worktree Completed in {:?}",
|
||||
t0.elapsed().as_millis()
|
||||
);
|
||||
}
|
||||
})
|
||||
.detach();
|
||||
|
||||
// let mut pending_files: Vec<(PathBuf, ((i64, PathBuf, Arc<Language>, SystemTime), SystemTime))> = vec![];
|
||||
this.update(&mut cx, |this, cx| {
|
||||
// The below is managing for updated on save
|
||||
// Currently each time a file is saved, this code is run, and for all the files that were changed, if the current time is
|
||||
// greater than the previous embedded time by the REINDEXING_DELAY variable, we will send the file off to be indexed.
|
||||
let _subscription = cx.subscribe(&project, |this, project, event, cx| {
|
||||
if let project::Event::WorktreeUpdatedEntries(worktree_id, changes) = event {
|
||||
this.project_entries_changed(project, changes.clone(), cx, worktree_id);
|
||||
}
|
||||
});
|
||||
|
||||
this.projects.insert(
|
||||
project.downgrade(),
|
||||
ProjectState {
|
||||
pending_files: HashMap::new(),
|
||||
worktree_db_ids: db_ids_by_worktree_id.into_iter().collect(),
|
||||
_subscription,
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
anyhow::Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
pub fn search(
|
||||
&mut self,
|
||||
project: ModelHandle<Project>,
|
||||
phrase: String,
|
||||
limit: usize,
|
||||
cx: &mut ModelContext<Self>,
|
||||
) -> Task<Result<Vec<SearchResult>>> {
|
||||
let project_state = if let Some(state) = self.projects.get(&project.downgrade()) {
|
||||
state
|
||||
} else {
|
||||
return Task::ready(Err(anyhow!("project not added")));
|
||||
};
|
||||
|
||||
let worktree_db_ids = project
|
||||
.read(cx)
|
||||
.worktrees(cx)
|
||||
.filter_map(|worktree| {
|
||||
let worktree_id = worktree.read(cx).id();
|
||||
project_state.db_id_for_worktree_id(worktree_id)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let embedding_provider = self.embedding_provider.clone();
|
||||
let database_url = self.database_url.clone();
|
||||
cx.spawn(|this, cx| async move {
|
||||
let documents = cx
|
||||
.background()
|
||||
.spawn(async move {
|
||||
let database = VectorDatabase::new(database_url.to_string_lossy().into())?;
|
||||
|
||||
let phrase_embedding = embedding_provider
|
||||
.embed_batch(vec![&phrase])
|
||||
.await?
|
||||
.into_iter()
|
||||
.next()
|
||||
.unwrap();
|
||||
|
||||
database.top_k_search(&worktree_db_ids, &phrase_embedding, limit)
|
||||
})
|
||||
.await?;
|
||||
|
||||
this.read_with(&cx, |this, _| {
|
||||
let project_state = if let Some(state) = this.projects.get(&project.downgrade()) {
|
||||
state
|
||||
} else {
|
||||
return Err(anyhow!("project not added"));
|
||||
};
|
||||
|
||||
Ok(documents
|
||||
.into_iter()
|
||||
.filter_map(|(worktree_db_id, file_path, offset, name)| {
|
||||
let worktree_id = project_state.worktree_id_for_db_id(worktree_db_id)?;
|
||||
Some(SearchResult {
|
||||
worktree_id,
|
||||
name,
|
||||
offset,
|
||||
file_path,
|
||||
})
|
||||
})
|
||||
.collect())
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn project_entries_changed(
|
||||
&mut self,
|
||||
project: ModelHandle<Project>,
|
||||
changes: Arc<[(Arc<Path>, ProjectEntryId, PathChange)]>,
|
||||
cx: &mut ModelContext<'_, VectorStore>,
|
||||
worktree_id: &WorktreeId,
|
||||
) -> Option<()> {
|
||||
let reindexing_delay = settings::get::<VectorStoreSettings>(cx).reindexing_delay_seconds;
|
||||
|
||||
let worktree = project
|
||||
.read(cx)
|
||||
.worktree_for_id(worktree_id.clone(), cx)?
|
||||
.read(cx)
|
||||
.snapshot();
|
||||
|
||||
let worktree_db_id = self
|
||||
.projects
|
||||
.get(&project.downgrade())?
|
||||
.db_id_for_worktree_id(worktree.id())?;
|
||||
let file_mtimes = self.get_file_mtimes(worktree_db_id);
|
||||
|
||||
let language_registry = self.language_registry.clone();
|
||||
|
||||
cx.spawn(|this, mut cx| async move {
|
||||
let file_mtimes = file_mtimes.await.log_err()?;
|
||||
|
||||
for change in changes.into_iter() {
|
||||
let change_path = change.0.clone();
|
||||
let absolute_path = worktree.absolutize(&change_path);
|
||||
|
||||
// Skip if git ignored or symlink
|
||||
if let Some(entry) = worktree.entry_for_id(change.1) {
|
||||
if entry.is_ignored || entry.is_symlink || entry.is_external {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
match change.2 {
|
||||
PathChange::Removed => this.update(&mut cx, |this, _| {
|
||||
this.db_update_tx
|
||||
.try_send(DbOperation::Delete {
|
||||
worktree_id: worktree_db_id,
|
||||
path: absolute_path,
|
||||
})
|
||||
.unwrap();
|
||||
}),
|
||||
_ => {
|
||||
if let Ok(language) = language_registry
|
||||
.language_for_file(&change_path.to_path_buf(), None)
|
||||
.await
|
||||
{
|
||||
if language
|
||||
.grammar()
|
||||
.and_then(|grammar| grammar.embedding_config.as_ref())
|
||||
.is_none()
|
||||
{
|
||||
continue;
|
||||
}
|
||||
|
||||
let modified_time =
|
||||
change_path.metadata().log_err()?.modified().log_err()?;
|
||||
|
||||
let existing_time = file_mtimes.get(&change_path.to_path_buf());
|
||||
let already_stored = existing_time
|
||||
.map_or(false, |existing_time| &modified_time != existing_time);
|
||||
|
||||
if !already_stored {
|
||||
this.update(&mut cx, |this, _| {
|
||||
let reindex_time = modified_time
|
||||
+ Duration::from_secs(reindexing_delay as u64);
|
||||
|
||||
let project_state =
|
||||
this.projects.get_mut(&project.downgrade())?;
|
||||
project_state.update_pending_files(
|
||||
PendingFile {
|
||||
relative_path: change_path.to_path_buf(),
|
||||
absolute_path,
|
||||
modified_time,
|
||||
worktree_db_id,
|
||||
language: language.clone(),
|
||||
},
|
||||
reindex_time,
|
||||
);
|
||||
|
||||
for file in project_state.get_outstanding_files() {
|
||||
this.parsing_files_tx.try_send(file).unwrap();
|
||||
}
|
||||
Some(())
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Some(())
|
||||
})
|
||||
.detach();
|
||||
|
||||
Some(())
|
||||
}
|
||||
}
|
||||
|
||||
impl Entity for VectorStore {
|
||||
type Event = ();
|
||||
}
|
||||
@@ -1,161 +0,0 @@
|
||||
use crate::{
|
||||
db::dot, embedding::EmbeddingProvider, vector_store_settings::VectorStoreSettings, VectorStore,
|
||||
};
|
||||
use anyhow::Result;
|
||||
use async_trait::async_trait;
|
||||
use gpui::{Task, TestAppContext};
|
||||
use language::{Language, LanguageConfig, LanguageRegistry};
|
||||
use project::{project_settings::ProjectSettings, FakeFs, Project};
|
||||
use rand::{rngs::StdRng, Rng};
|
||||
use serde_json::json;
|
||||
use settings::SettingsStore;
|
||||
use std::sync::Arc;
|
||||
use unindent::Unindent;
|
||||
|
||||
#[gpui::test]
|
||||
async fn test_vector_store(cx: &mut TestAppContext) {
|
||||
cx.update(|cx| {
|
||||
cx.set_global(SettingsStore::test(cx));
|
||||
settings::register::<VectorStoreSettings>(cx);
|
||||
settings::register::<ProjectSettings>(cx);
|
||||
});
|
||||
|
||||
let fs = FakeFs::new(cx.background());
|
||||
fs.insert_tree(
|
||||
"/the-root",
|
||||
json!({
|
||||
"src": {
|
||||
"file1.rs": "
|
||||
fn aaa() {
|
||||
println!(\"aaaa!\");
|
||||
}
|
||||
|
||||
fn zzzzzzzzz() {
|
||||
println!(\"SLEEPING\");
|
||||
}
|
||||
".unindent(),
|
||||
"file2.rs": "
|
||||
fn bbb() {
|
||||
println!(\"bbbb!\");
|
||||
}
|
||||
".unindent(),
|
||||
}
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
let languages = Arc::new(LanguageRegistry::new(Task::ready(())));
|
||||
let rust_language = Arc::new(
|
||||
Language::new(
|
||||
LanguageConfig {
|
||||
name: "Rust".into(),
|
||||
path_suffixes: vec!["rs".into()],
|
||||
..Default::default()
|
||||
},
|
||||
Some(tree_sitter_rust::language()),
|
||||
)
|
||||
.with_embedding_query(
|
||||
r#"
|
||||
(function_item
|
||||
name: (identifier) @name
|
||||
body: (block)) @item
|
||||
"#,
|
||||
)
|
||||
.unwrap(),
|
||||
);
|
||||
languages.add(rust_language);
|
||||
|
||||
let db_dir = tempdir::TempDir::new("vector-store").unwrap();
|
||||
let db_path = db_dir.path().join("db.sqlite");
|
||||
|
||||
let store = VectorStore::new(
|
||||
fs.clone(),
|
||||
db_path,
|
||||
Arc::new(FakeEmbeddingProvider),
|
||||
languages,
|
||||
cx.to_async(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let project = Project::test(fs, ["/the-root".as_ref()], cx).await;
|
||||
let worktree_id = project.read_with(cx, |project, cx| {
|
||||
project.worktrees(cx).next().unwrap().read(cx).id()
|
||||
});
|
||||
store
|
||||
.update(cx, |store, cx| store.add_project(project.clone(), cx))
|
||||
.await
|
||||
.unwrap();
|
||||
cx.foreground().run_until_parked();
|
||||
|
||||
let search_results = store
|
||||
.update(cx, |store, cx| {
|
||||
store.search(project.clone(), "aaaa".to_string(), 5, cx)
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(search_results[0].offset, 0);
|
||||
assert_eq!(search_results[0].name, "aaa");
|
||||
assert_eq!(search_results[0].worktree_id, worktree_id);
|
||||
}
|
||||
|
||||
#[gpui::test]
|
||||
fn test_dot_product(mut rng: StdRng) {
|
||||
assert_eq!(dot(&[1., 0., 0., 0., 0.], &[0., 1., 0., 0., 0.]), 0.);
|
||||
assert_eq!(dot(&[2., 0., 0., 0., 0.], &[3., 1., 0., 0., 0.]), 6.);
|
||||
|
||||
for _ in 0..100 {
|
||||
let size = 1536;
|
||||
let mut a = vec![0.; size];
|
||||
let mut b = vec![0.; size];
|
||||
for (a, b) in a.iter_mut().zip(b.iter_mut()) {
|
||||
*a = rng.gen();
|
||||
*b = rng.gen();
|
||||
}
|
||||
|
||||
assert_eq!(
|
||||
round_to_decimals(dot(&a, &b), 1),
|
||||
round_to_decimals(reference_dot(&a, &b), 1)
|
||||
);
|
||||
}
|
||||
|
||||
fn round_to_decimals(n: f32, decimal_places: i32) -> f32 {
|
||||
let factor = (10.0 as f32).powi(decimal_places);
|
||||
(n * factor).round() / factor
|
||||
}
|
||||
|
||||
fn reference_dot(a: &[f32], b: &[f32]) -> f32 {
|
||||
a.iter().zip(b.iter()).map(|(a, b)| a * b).sum()
|
||||
}
|
||||
}
|
||||
|
||||
struct FakeEmbeddingProvider;
|
||||
|
||||
#[async_trait]
|
||||
impl EmbeddingProvider for FakeEmbeddingProvider {
|
||||
async fn embed_batch(&self, spans: Vec<&str>) -> Result<Vec<Vec<f32>>> {
|
||||
Ok(spans
|
||||
.iter()
|
||||
.map(|span| {
|
||||
let mut result = vec![1.0; 26];
|
||||
for letter in span.chars() {
|
||||
let letter = letter.to_ascii_lowercase();
|
||||
if letter as u32 >= 'a' as u32 {
|
||||
let ix = (letter as u32) - ('a' as u32);
|
||||
if ix < 26 {
|
||||
result[ix as usize] += 1.0;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let norm = result.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
for x in &mut result {
|
||||
*x /= norm;
|
||||
}
|
||||
|
||||
result
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
}
|
||||
@@ -64,7 +64,7 @@ terminal_view = { path = "../terminal_view" }
|
||||
theme = { path = "../theme" }
|
||||
theme_selector = { path = "../theme_selector" }
|
||||
util = { path = "../util" }
|
||||
vector_store = { path = "../vector_store" }
|
||||
semantic_index = { path = "../semantic_index" }
|
||||
vim = { path = "../vim" }
|
||||
workspace = { path = "../workspace" }
|
||||
welcome = { path = "../welcome" }
|
||||
|
||||
@@ -0,0 +1,43 @@
|
||||
(
|
||||
(comment)* @context
|
||||
.
|
||||
(declaration
|
||||
declarator: [
|
||||
(function_declarator
|
||||
declarator: (_) @name)
|
||||
(pointer_declarator
|
||||
"*" @name
|
||||
declarator: (function_declarator
|
||||
declarator: (_) @name))
|
||||
(pointer_declarator
|
||||
"*" @name
|
||||
declarator: (pointer_declarator
|
||||
"*" @name
|
||||
declarator: (function_declarator
|
||||
declarator: (_) @name)))
|
||||
]
|
||||
) @item
|
||||
)
|
||||
|
||||
(
|
||||
(comment)* @context
|
||||
.
|
||||
(function_definition
|
||||
declarator: [
|
||||
(function_declarator
|
||||
declarator: (_) @name
|
||||
)
|
||||
(pointer_declarator
|
||||
"*" @name
|
||||
declarator: (function_declarator
|
||||
declarator: (_) @name
|
||||
))
|
||||
(pointer_declarator
|
||||
"*" @name
|
||||
declarator: (pointer_declarator
|
||||
"*" @name
|
||||
declarator: (function_declarator
|
||||
declarator: (_) @name)))
|
||||
]
|
||||
) @item
|
||||
)
|
||||
@@ -0,0 +1,61 @@
|
||||
(
|
||||
(comment)* @context
|
||||
.
|
||||
(function_definition
|
||||
(type_qualifier)? @name
|
||||
type: (_)? @name
|
||||
declarator: [
|
||||
(function_declarator
|
||||
declarator: (_) @name)
|
||||
(pointer_declarator
|
||||
"*" @name
|
||||
declarator: (function_declarator
|
||||
declarator: (_) @name))
|
||||
(pointer_declarator
|
||||
"*" @name
|
||||
declarator: (pointer_declarator
|
||||
"*" @name
|
||||
declarator: (function_declarator
|
||||
declarator: (_) @name)))
|
||||
(reference_declarator
|
||||
["&" "&&"] @name
|
||||
(function_declarator
|
||||
declarator: (_) @name))
|
||||
]
|
||||
(type_qualifier)? @name) @item
|
||||
)
|
||||
|
||||
(
|
||||
(comment)* @context
|
||||
.
|
||||
(template_declaration
|
||||
(class_specifier
|
||||
"class" @name
|
||||
name: (_) @name)
|
||||
) @item
|
||||
)
|
||||
|
||||
(
|
||||
(comment)* @context
|
||||
.
|
||||
(class_specifier
|
||||
"class" @name
|
||||
name: (_) @name) @item
|
||||
)
|
||||
|
||||
(
|
||||
(comment)* @context
|
||||
.
|
||||
(enum_specifier
|
||||
"enum" @name
|
||||
name: (_) @name) @item
|
||||
)
|
||||
|
||||
(
|
||||
(comment)* @context
|
||||
.
|
||||
(declaration
|
||||
type: (struct_specifier
|
||||
"struct" @name)
|
||||
declarator: (_) @name) @item
|
||||
)
|
||||
@@ -0,0 +1,27 @@
|
||||
(
|
||||
(unary_operator
|
||||
operator: "@"
|
||||
operand: (call
|
||||
target: (identifier) @unary
|
||||
(#match? @unary "^(doc)$"))
|
||||
) @context
|
||||
.
|
||||
(call
|
||||
target: (identifier) @name
|
||||
(arguments
|
||||
[
|
||||
(identifier) @name
|
||||
(call
|
||||
target: (identifier) @name)
|
||||
(binary_operator
|
||||
left: (call
|
||||
target: (identifier) @name)
|
||||
operator: "when")
|
||||
])
|
||||
(#match? @name "^(def|defp|defdelegate|defguard|defguardp|defmacro|defmacrop|defn|defnp)$")) @item
|
||||
)
|
||||
|
||||
(call
|
||||
target: (identifier) @name
|
||||
(arguments (alias) @name)
|
||||
(#match? @name "^(defmodule|defprotocol)$")) @item
|
||||
@@ -0,0 +1,24 @@
|
||||
(
|
||||
(comment)* @context
|
||||
.
|
||||
(type_declaration
|
||||
(type_spec
|
||||
name: (_) @name)
|
||||
) @item
|
||||
)
|
||||
|
||||
(
|
||||
(comment)* @context
|
||||
.
|
||||
(function_declaration
|
||||
name: (_) @name
|
||||
) @item
|
||||
)
|
||||
|
||||
(
|
||||
(comment)* @context
|
||||
.
|
||||
(method_declaration
|
||||
name: (_) @name
|
||||
) @item
|
||||
)
|
||||
@@ -1,56 +1,71 @@
|
||||
; (internal_module
|
||||
; "namespace" @context
|
||||
; name: (_) @name) @item
|
||||
|
||||
(enum_declaration
|
||||
"enum" @context
|
||||
name: (_) @name) @item
|
||||
|
||||
(function_declaration
|
||||
"async"? @context
|
||||
"function" @context
|
||||
name: (_) @name) @item
|
||||
|
||||
(interface_declaration
|
||||
"interface" @context
|
||||
name: (_) @name) @item
|
||||
|
||||
; (program
|
||||
; (export_statement
|
||||
; (lexical_declaration
|
||||
; ["let" "const"] @context
|
||||
; (variable_declarator
|
||||
; name: (_) @name) @item)))
|
||||
|
||||
(program
|
||||
(lexical_declaration
|
||||
["let" "const"] @context
|
||||
(variable_declarator
|
||||
name: (_) @name) @item))
|
||||
|
||||
(class_declaration
|
||||
"class" @context
|
||||
name: (_) @name) @item
|
||||
|
||||
(method_definition
|
||||
(
|
||||
(comment)* @context
|
||||
.
|
||||
[
|
||||
"get"
|
||||
"set"
|
||||
"async"
|
||||
"*"
|
||||
"readonly"
|
||||
"static"
|
||||
(override_modifier)
|
||||
(accessibility_modifier)
|
||||
]* @context
|
||||
name: (_) @name) @item
|
||||
(export_statement
|
||||
(function_declaration
|
||||
"async"? @name
|
||||
"function" @name
|
||||
name: (_) @name))
|
||||
(function_declaration
|
||||
"async"? @name
|
||||
"function" @name
|
||||
name: (_) @name)
|
||||
] @item
|
||||
)
|
||||
|
||||
; (public_field_definition
|
||||
; [
|
||||
; "declare"
|
||||
; "readonly"
|
||||
; "abstract"
|
||||
; "static"
|
||||
; (accessibility_modifier)
|
||||
; ]* @context
|
||||
; name: (_) @name) @item
|
||||
(
|
||||
(comment)* @context
|
||||
.
|
||||
[
|
||||
(export_statement
|
||||
(class_declaration
|
||||
"class" @name
|
||||
name: (_) @name))
|
||||
(class_declaration
|
||||
"class" @name
|
||||
name: (_) @name)
|
||||
] @item
|
||||
)
|
||||
|
||||
(
|
||||
(comment)* @context
|
||||
.
|
||||
[
|
||||
(export_statement
|
||||
(interface_declaration
|
||||
"interface" @name
|
||||
name: (_) @name))
|
||||
(interface_declaration
|
||||
"interface" @name
|
||||
name: (_) @name)
|
||||
] @item
|
||||
)
|
||||
|
||||
(
|
||||
(comment)* @context
|
||||
.
|
||||
[
|
||||
(export_statement
|
||||
(enum_declaration
|
||||
"enum" @name
|
||||
name: (_) @name))
|
||||
(enum_declaration
|
||||
"enum" @name
|
||||
name: (_) @name)
|
||||
] @item
|
||||
)
|
||||
|
||||
(
|
||||
(comment)* @context
|
||||
.
|
||||
(method_definition
|
||||
[
|
||||
"get"
|
||||
"set"
|
||||
"async"
|
||||
"*"
|
||||
"static"
|
||||
]* @name
|
||||
name: (_) @name) @item
|
||||
)
|
||||
|
||||
@@ -0,0 +1,14 @@
|
||||
; Only produce one embedding for the entire file.
|
||||
(document) @item
|
||||
|
||||
; Collapse arrays, except for the first object.
|
||||
(array
|
||||
"[" @keep
|
||||
.
|
||||
(object)? @keep
|
||||
"]" @keep) @collapse
|
||||
|
||||
; Collapse string values (but not keys).
|
||||
(pair value: (string
|
||||
"\"" @keep
|
||||
"\"" @keep) @collapse)
|
||||
@@ -10,3 +10,4 @@ brackets = [
|
||||
{ start = "\"", end = "\"", close = true, newline = false, not_in = ["string"] },
|
||||
{ start = "/*", end = " */", close = true, newline = false, not_in = ["string", "comment"] },
|
||||
]
|
||||
collapsed_placeholder = " /* ... */ "
|
||||
|
||||
@@ -1,36 +1,28 @@
|
||||
(struct_item
|
||||
(visibility_modifier)? @context
|
||||
"struct" @context
|
||||
name: (_) @name) @item
|
||||
(
|
||||
[(line_comment) (attribute_item)]* @context
|
||||
.
|
||||
[
|
||||
(struct_item
|
||||
name: (_) @name)
|
||||
|
||||
(enum_item
|
||||
(visibility_modifier)? @context
|
||||
"enum" @context
|
||||
name: (_) @name) @item
|
||||
(enum_item
|
||||
name: (_) @name)
|
||||
|
||||
(impl_item
|
||||
"impl" @context
|
||||
trait: (_)? @name
|
||||
"for"? @context
|
||||
type: (_) @name) @item
|
||||
(impl_item
|
||||
trait: (_)? @name
|
||||
"for"? @name
|
||||
type: (_) @name)
|
||||
|
||||
(trait_item
|
||||
(visibility_modifier)? @context
|
||||
"trait" @context
|
||||
name: (_) @name) @item
|
||||
(trait_item
|
||||
name: (_) @name)
|
||||
|
||||
(function_item
|
||||
(visibility_modifier)? @context
|
||||
(function_modifiers)? @context
|
||||
"fn" @context
|
||||
name: (_) @name) @item
|
||||
(function_item
|
||||
name: (_) @name
|
||||
body: (block
|
||||
"{" @keep
|
||||
"}" @keep) @collapse)
|
||||
|
||||
(function_signature_item
|
||||
(visibility_modifier)? @context
|
||||
(function_modifiers)? @context
|
||||
"fn" @context
|
||||
name: (_) @name) @item
|
||||
|
||||
(macro_definition
|
||||
. "macro_rules!" @context
|
||||
name: (_) @name) @item
|
||||
(macro_definition
|
||||
name: (_) @name)
|
||||
] @item
|
||||
)
|
||||
|
||||
@@ -1,35 +1,85 @@
|
||||
(enum_declaration
|
||||
"enum" @context
|
||||
name: (_) @name) @item
|
||||
|
||||
(function_declaration
|
||||
"async"? @context
|
||||
"function" @context
|
||||
name: (_) @name) @item
|
||||
|
||||
(interface_declaration
|
||||
"interface" @context
|
||||
name: (_) @name) @item
|
||||
|
||||
(program
|
||||
(lexical_declaration
|
||||
["let" "const"] @context
|
||||
(variable_declarator
|
||||
name: (_) @name) @item))
|
||||
|
||||
(class_declaration
|
||||
"class" @context
|
||||
name: (_) @name) @item
|
||||
|
||||
(method_definition
|
||||
(
|
||||
(comment)* @context
|
||||
.
|
||||
[
|
||||
"get"
|
||||
"set"
|
||||
"async"
|
||||
"*"
|
||||
"readonly"
|
||||
"static"
|
||||
(override_modifier)
|
||||
(accessibility_modifier)
|
||||
]* @context
|
||||
name: (_) @name) @item
|
||||
(export_statement
|
||||
(function_declaration
|
||||
"async"? @name
|
||||
"function" @name
|
||||
name: (_) @name))
|
||||
(function_declaration
|
||||
"async"? @name
|
||||
"function" @name
|
||||
name: (_) @name)
|
||||
] @item
|
||||
)
|
||||
|
||||
(
|
||||
(comment)* @context
|
||||
.
|
||||
[
|
||||
(export_statement
|
||||
(class_declaration
|
||||
"class" @name
|
||||
name: (_) @name))
|
||||
(class_declaration
|
||||
"class" @name
|
||||
name: (_) @name)
|
||||
] @item
|
||||
)
|
||||
|
||||
(
|
||||
(comment)* @context
|
||||
.
|
||||
[
|
||||
(export_statement
|
||||
(interface_declaration
|
||||
"interface" @name
|
||||
name: (_) @name))
|
||||
(interface_declaration
|
||||
"interface" @name
|
||||
name: (_) @name)
|
||||
] @item
|
||||
)
|
||||
|
||||
(
|
||||
(comment)* @context
|
||||
.
|
||||
[
|
||||
(export_statement
|
||||
(enum_declaration
|
||||
"enum" @name
|
||||
name: (_) @name))
|
||||
(enum_declaration
|
||||
"enum" @name
|
||||
name: (_) @name)
|
||||
] @item
|
||||
)
|
||||
|
||||
(
|
||||
(comment)* @context
|
||||
.
|
||||
[
|
||||
(export_statement
|
||||
(type_alias_declaration
|
||||
"type" @name
|
||||
name: (_) @name))
|
||||
(type_alias_declaration
|
||||
"type" @name
|
||||
name: (_) @name)
|
||||
] @item
|
||||
)
|
||||
|
||||
(
|
||||
(comment)* @context
|
||||
.
|
||||
(method_definition
|
||||
[
|
||||
"get"
|
||||
"set"
|
||||
"async"
|
||||
"*"
|
||||
"static"
|
||||
]* @name
|
||||
name: (_) @name) @item
|
||||
)
|
||||
|
||||
@@ -1,59 +1,85 @@
|
||||
; (internal_module
|
||||
; "namespace" @context
|
||||
; name: (_) @name) @item
|
||||
|
||||
(enum_declaration
|
||||
"enum" @context
|
||||
name: (_) @name) @item
|
||||
|
||||
; (type_alias_declaration
|
||||
; "type" @context
|
||||
; name: (_) @name) @item
|
||||
|
||||
(function_declaration
|
||||
"async"? @context
|
||||
"function" @context
|
||||
name: (_) @name) @item
|
||||
|
||||
(interface_declaration
|
||||
"interface" @context
|
||||
name: (_) @name) @item
|
||||
|
||||
; (export_statement
|
||||
; (lexical_declaration
|
||||
; ["let" "const"] @context
|
||||
; (variable_declarator
|
||||
; name: (_) @name) @item))
|
||||
|
||||
(program
|
||||
(lexical_declaration
|
||||
["let" "const"] @context
|
||||
(variable_declarator
|
||||
name: (_) @name) @item))
|
||||
|
||||
(class_declaration
|
||||
"class" @context
|
||||
name: (_) @name) @item
|
||||
|
||||
(method_definition
|
||||
(
|
||||
(comment)* @context
|
||||
.
|
||||
[
|
||||
"get"
|
||||
"set"
|
||||
"async"
|
||||
"*"
|
||||
"readonly"
|
||||
"static"
|
||||
(override_modifier)
|
||||
(accessibility_modifier)
|
||||
]* @context
|
||||
name: (_) @name) @item
|
||||
(export_statement
|
||||
(function_declaration
|
||||
"async"? @name
|
||||
"function" @name
|
||||
name: (_) @name))
|
||||
(function_declaration
|
||||
"async"? @name
|
||||
"function" @name
|
||||
name: (_) @name)
|
||||
] @item
|
||||
)
|
||||
|
||||
; (public_field_definition
|
||||
; [
|
||||
; "declare"
|
||||
; "readonly"
|
||||
; "abstract"
|
||||
; "static"
|
||||
; (accessibility_modifier)
|
||||
; ]* @context
|
||||
; name: (_) @name) @item
|
||||
(
|
||||
(comment)* @context
|
||||
.
|
||||
[
|
||||
(export_statement
|
||||
(class_declaration
|
||||
"class" @name
|
||||
name: (_) @name))
|
||||
(class_declaration
|
||||
"class" @name
|
||||
name: (_) @name)
|
||||
] @item
|
||||
)
|
||||
|
||||
(
|
||||
(comment)* @context
|
||||
.
|
||||
[
|
||||
(export_statement
|
||||
(interface_declaration
|
||||
"interface" @name
|
||||
name: (_) @name))
|
||||
(interface_declaration
|
||||
"interface" @name
|
||||
name: (_) @name)
|
||||
] @item
|
||||
)
|
||||
|
||||
(
|
||||
(comment)* @context
|
||||
.
|
||||
[
|
||||
(export_statement
|
||||
(enum_declaration
|
||||
"enum" @name
|
||||
name: (_) @name))
|
||||
(enum_declaration
|
||||
"enum" @name
|
||||
name: (_) @name)
|
||||
] @item
|
||||
)
|
||||
|
||||
(
|
||||
(comment)* @context
|
||||
.
|
||||
[
|
||||
(export_statement
|
||||
(type_alias_declaration
|
||||
"type" @name
|
||||
name: (_) @name))
|
||||
(type_alias_declaration
|
||||
"type" @name
|
||||
name: (_) @name)
|
||||
] @item
|
||||
)
|
||||
|
||||
(
|
||||
(comment)* @context
|
||||
.
|
||||
(method_definition
|
||||
[
|
||||
"get"
|
||||
"set"
|
||||
"async"
|
||||
"*"
|
||||
"static"
|
||||
]* @name
|
||||
name: (_) @name) @item
|
||||
)
|
||||
|
||||
@@ -157,7 +157,7 @@ fn main() {
|
||||
project_panel::init(Assets, cx);
|
||||
diagnostics::init(cx);
|
||||
search::init(cx);
|
||||
vector_store::init(fs.clone(), http.clone(), languages.clone(), cx);
|
||||
semantic_index::init(fs.clone(), http.clone(), languages.clone(), cx);
|
||||
vim::init(cx);
|
||||
terminal_view::init(cx);
|
||||
copilot::init(http.clone(), node_runtime, cx);
|
||||
|
||||
Reference in New Issue
Block a user