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

updated vector store indexing to only use languages with an embedding.scm treesitter query

Browse source 
Co-authored-by: maxbrunsfeld <max@zed.dev>
This commit is contained in:
KCaverly 2023-06-30 16:14:11 -04:00
parent 0db0876289
commit 36907bb4dc
5 changed files with 98 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

44
crates/language/src/language.rs
View file

@ -350,6 +350,7 @@ pub struct LanguageQueries {
pub brackets: Option<Cow<'static, str>>,
pub indents: Option<Cow<'static, str>>,
pub outline: Option<Cow<'static, str>>,
pub embedding: Option<Cow<'static, str>>,
pub injections: Option<Cow<'static, str>>,
pub overrides: Option<Cow<'static, str>>,
}
@ -495,6 +496,7 @@ pub struct Grammar {
pub(crate) brackets_config: Option<BracketConfig>,
pub(crate) indents_config: Option<IndentConfig>,
pub outline_config: Option<OutlineConfig>,
pub embedding_config: Option<EmbeddingConfig>,
pub(crate) injection_config: Option<InjectionConfig>,
pub(crate) override_config: Option<OverrideConfig>,
pub(crate) highlight_map: Mutex<HighlightMap>,
@ -516,6 +518,15 @@ pub struct OutlineConfig {
pub extra_context_capture_ix: Option<u32>,
}
#[derive(Debug)]
pub struct EmbeddingConfig {
pub query: Query,
pub item_capture_ix: u32,
pub name_capture_ix: u32,
pub context_capture_ix: Option<u32>,
pub extra_context_capture_ix: Option<u32>,
}
struct InjectionConfig {
query: Query,
content_capture_ix: u32,
@ -1145,6 +1156,7 @@ impl Language {
highlights_query: None,
brackets_config: None,
outline_config: None,
embedding_config: None,
indents_config: None,
injection_config: None,
override_config: None,
@ -1181,6 +1193,9 @@ impl Language {
if let Some(query) = queries.outline {
self = self.with_outline_query(query.as_ref())?;
}
if let Some(query) = queries.embedding {
self = self.with_embedding_query(query.as_ref())?;
}
if let Some(query) = queries.injections {
self = self.with_injection_query(query.as_ref())?;
}
@ -1189,6 +1204,7 @@ impl Language {
}
Ok(self)
}
pub fn with_highlights_query(mut self, source: &str) -> Result<Self> {
let grammar = self.grammar_mut();
grammar.highlights_query = Some(Query::new(grammar.ts_language, source)?);
@ -1223,6 +1239,34 @@ impl Language {
Ok(self)
}
pub fn with_embedding_query(mut self, source: &str) -> Result<Self> {
let grammar = self.grammar_mut();
let query = Query::new(grammar.ts_language, source)?;
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;
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),
],
);
if let Some((item_capture_ix, name_capture_ix)) = item_capture_ix.zip(name_capture_ix) {
grammar.embedding_config = Some(EmbeddingConfig {
query,
item_capture_ix,
name_capture_ix,
context_capture_ix,
extra_context_capture_ix,
});
}
Ok(self)
}
pub fn with_brackets_query(mut self, source: &str) -> Result<Self> {
let grammar = self.grammar_mut();
let query = Query::new(grammar.ts_language, source)?;