add embedding query for json with nested arrays and strings

Co-authored-by: maxbrunsfeld <max@zed.dev>
This commit is contained in:
KCaverly 2023-07-19 16:52:44 -04:00
parent 9809ec3d70
commit efe973ebe2
7 changed files with 189 additions and 59 deletions

View file

@ -170,9 +170,7 @@ async fn test_code_context_retrieval_rust() {
"
.unindent();
let documents = retriever
.parse_file(Path::new("foo.rs"), &text, language)
.unwrap();
let documents = retriever.parse_file(&text, language).unwrap();
assert_documents_eq(
&documents,
@ -229,6 +227,76 @@ async fn test_code_context_retrieval_rust() {
);
}
#[gpui::test]
async fn test_code_context_retrieval_json() {
let language = json_lang();
let mut retriever = CodeContextRetriever::new();
let text = r#"
{
"array": [1, 2, 3, 4],
"string": "abcdefg",
"nested_object": {
"array_2": [5, 6, 7, 8],
"string_2": "hijklmnop",
"boolean": true,
"none": null
}
}
"#
.unindent();
let documents = retriever.parse_file(&text, language.clone()).unwrap();
assert_documents_eq(
&documents,
&[(
r#"
{
"array": [],
"string": "",
"nested_object": {
"array_2": [],
"string_2": "",
"boolean": true,
"none": null
}
}"#
.unindent(),
text.find("{").unwrap(),
)],
);
let text = r#"
[
{
"name": "somebody",
"age": 42
},
{
"name": "somebody else",
"age": 43
}
]
"#
.unindent();
let documents = retriever.parse_file(&text, language.clone()).unwrap();
assert_documents_eq(
&documents,
&[(
r#"
[{
"name": "",
"age": 42
}]"#
.unindent(),
text.find("[").unwrap(),
)],
);
}
fn assert_documents_eq(
documents: &[Document],
expected_contents_and_start_offsets: &[(String, usize)],
@ -913,6 +981,35 @@ fn rust_lang() -> Arc<Language> {
)
}
fn json_lang() -> Arc<Language> {
Arc::new(
Language::new(
LanguageConfig {
name: "JSON".into(),
path_suffixes: vec!["json".into()],
..Default::default()
},
Some(tree_sitter_json::language()),
)
.with_embedding_query(
r#"
(document) @item
(array
"[" @keep
.
(object)? @keep
"]" @keep) @collapse
(pair value: (string
"\"" @keep
"\"" @keep) @collapse)
"#,
)
.unwrap(),
)
}
fn toml_lang() -> Arc<Language> {
Arc::new(Language::new(
LanguageConfig {