Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

add glob filtering functionality to semantic search

Browse source
This commit is contained in:
KCaverly 2023-07-20 13:46:27 -04:00
parent efe973ebe2
commit e02d6bc0d4
6 changed files with 149 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

1
crates/semantic_index/Cargo.toml
View file

@ -37,6 +37,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"] }

39
crates/semantic_index/src/db.rs
View file

@ -1,5 +1,6 @@
use crate::{parsing::Document, SEMANTIC_INDEX_VERSION};
use anyhow::{anyhow, Context, Result};
use globset::{Glob, GlobMatcher};
use project::Fs;
use rpc::proto::Timestamp;
use rusqlite::{
@ -252,18 +253,30 @@ impl VectorDatabase {
worktree_ids: &[i64],
query_embedding: &Vec<f32>,
limit: usize,
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))
self.for_each_document(&worktree_ids, |relative_path, id, embedding| {
if (include_globs.is_empty()
|| include_globs
.iter()
.any(|include_glob| include_glob.is_match(relative_path.clone())))
&& (exclude_globs.is_empty()
|| !exclude_globs
.iter()
.any(|exclude_glob| exclude_glob.is_match(relative_path.clone())))
{
Ok(ix) => ix,
Err(ix) => ix,
};
results.insert(ix, (id, similarity));
results.truncate(limit);
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<_>>();
@ -273,12 +286,12 @@ impl VectorDatabase {
fn for_each_document(
&self,
worktree_ids: &[i64],
mut f: impl FnMut(i64, Vec<f32>),
mut f: impl FnMut(String, i64, Vec<f32>),
) -> Result<()> {
let mut query_statement = self.db.prepare(
"
SELECT
documents.id, documents.embedding
files.relative_path, documents.id, documents.embedding
FROM
documents, files
WHERE
@ -289,10 +302,10 @@ impl VectorDatabase {
query_statement
.query_map(params![ids_to_sql(worktree_ids)], |row| {
Ok((row.get(0)?, row.get::<_, Embedding>(1)?))
Ok((row.get(0)?, row.get(1)?, row.get::<_, Embedding>(2)?))
})?
.filter_map(|row| row.ok())
.for_each(|(id, embedding)| f(id, embedding.0));
.for_each(|(relative_path, id, embedding)| f(relative_path, id, embedding.0));
Ok(())
}

13
crates/semantic_index/src/semantic_index.rs
View file

@ -11,6 +11,7 @@ use anyhow::{anyhow, Result};
use db::VectorDatabase;
use embedding::{EmbeddingProvider, OpenAIEmbeddings};
use futures::{channel::oneshot, Future};
use globset::{Glob, GlobMatcher};
use gpui::{AppContext, AsyncAppContext, Entity, ModelContext, ModelHandle, Task, WeakModelHandle};
use language::{Anchor, Buffer, Language, LanguageRegistry};
use parking_lot::Mutex;
@ -624,6 +625,8 @@ impl SemanticIndex {
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()) {
@ -657,12 +660,16 @@ impl SemanticIndex {
.next()
.unwrap();
database.top_k_search(&worktree_db_ids, &phrase_embedding, limit)
database.top_k_search(
&worktree_db_ids,
&phrase_embedding,
limit,
include_globs,
exclude_globs,
)
})
.await?;
dbg!(&documents);
let mut tasks = Vec::new();
let mut ranges = Vec::new();
let weak_project = project.downgrade();

57
crates/semantic_index/src/semantic_index_tests.rs
View file

@ -7,6 +7,7 @@ use crate::{
};
use anyhow::Result;
use async_trait::async_trait;
use globset::Glob;
use gpui::{Task, TestAppContext};
use language::{Language, LanguageConfig, LanguageRegistry, ToOffset};
use pretty_assertions::assert_eq;
@ -96,7 +97,7 @@ async fn test_semantic_index(cx: &mut TestAppContext) {
let search_results = store
.update(cx, |store, cx| {
store.search_project(project.clone(), "aaaa".to_string(), 5, cx)
store.search_project(project.clone(), "aaaa".to_string(), 5, vec![], vec![], cx)
})
.await
.unwrap();
@ -109,6 +110,60 @@ async fn test_semantic_index(cx: &mut TestAppContext) {
);
});
// Test Include Files Functonality
let include_files = vec![Glob::new("*.rs").unwrap().compile_matcher()];
let exclude_files = vec![Glob::new("*.rs").unwrap().compile_matcher()];
let search_results = store
.update(cx, |store, cx| {
store.search_project(
project.clone(),
"aaaa".to_string(),
5,
include_files,
vec![],
cx,
)
})
.await
.unwrap();
for res in &search_results {
res.buffer.read_with(cx, |buffer, _cx| {
assert!(buffer
.file()
.unwrap()
.path()
.to_str()
.unwrap()
.ends_with("rs"));
});
}
let search_results = store
.update(cx, |store, cx| {
store.search_project(
project.clone(),
"aaaa".to_string(),
5,
vec![],
exclude_files,
cx,
)
})
.await
.unwrap();
for res in &search_results {
res.buffer.read_with(cx, |buffer, _cx| {
assert!(!buffer
.file()
.unwrap()
.path()
.to_str()
.unwrap()
.ends_with("rs"));
});
}
fs.save(
"/the-root/src/file2.rs".as_ref(),
&"