feat: add rust prompt rewrite support (TabbyML#296)

javadeath · Jul 13, 2023 · be5fe0d · be5fe0d
1 parent 4388fd0
commit be5fe0d
Show file tree

Hide file tree

Showing 4 changed files with 43 additions and 14 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/tabby-scheduler/Cargo.toml b/crates/tabby-scheduler/Cargo.toml
@@ -20,6 +20,7 @@ serde = { workspace = true }
 serde-jsonlines = { workspace = true }
 file-rotate = "0.7.5"
 tree-sitter-python = "0.20.2"
+tree-sitter-rust = "0.20.3"
 
 [dev-dependencies]
 temp_testdir = "0.2"

diff --git a/crates/tabby-scheduler/src/dataset.rs b/crates/tabby-scheduler/src/dataset.rs
@@ -224,16 +224,29 @@ lazy_static! {
         map
     };
     static ref LANGUAGE_TAGS: HashMap<&'static str, TagsConfigurationSync> = {
-        HashMap::from([(
-            "python",
-            TagsConfigurationSync(
-                TagsConfiguration::new(
-                    tree_sitter_python::language(),
-                    tree_sitter_python::TAGGING_QUERY,
-                    "",
-                )
-                .unwrap(),
+        HashMap::from([
+            (
+                "python",
+                TagsConfigurationSync(
+                    TagsConfiguration::new(
+                        tree_sitter_python::language(),
+                        tree_sitter_python::TAGGING_QUERY,
+                        "",
+                    )
+                    .unwrap(),
+                ),
             ),
-        )])
+            (
+                "rust",
+                TagsConfigurationSync(
+                    TagsConfiguration::new(
+                        tree_sitter_rust::language(),
+                        tree_sitter_rust::TAGGING_QUERY,
+                        "",
+                    )
+                    .unwrap(),
+                ),
+            ),
+        ])
     };
 }
diff --git a/crates/tabby/src/serve/completions/prompt.rs b/crates/tabby/src/serve/completions/prompt.rs
@@ -155,8 +155,12 @@ fn collect_snippets(index: &IndexState, language: &str, text: &str) -> Vec<Strin
 }
 
 fn sanitize_text(text: &str) -> String {
-    let x = text.replace(|c: char| !c.is_ascii_digit() && !c.is_alphabetic(), " ");
-    let tokens: Vec<&str> = x.split(' ').collect();
+    // Only keep [a-zA-Z0-9-_]
+    let x = text.replace(
+        |c: char| !c.is_ascii_digit() && !c.is_alphabetic() && c != '_' && c != '-',
+        " ",
+    );
+    let tokens: Vec<&str> = x.split(' ').filter(|x| x.len() > 5).collect();
     tokens.join(" ")
 }
 
@@ -182,7 +186,7 @@ impl IndexState {
             .schema()
             .get_field("body")
             .ok_or(anyhow!("Index doesn't have required field"))?;
-        let query_parser = QueryParser::for_index(&index, vec![field_name]);
+        let query_parser = QueryParser::for_index(&index, vec![field_body]);
         Ok(Self {
             searcher: reader.searcher(),
             query_parser,
@@ -194,5 +198,5 @@ impl IndexState {
 
 lazy_static! {
     static ref LANGUAGE_LINE_COMMENT_CHAR: HashMap<&'static str, &'static str> =
-        HashMap::from([("python", "#")]);
+        HashMap::from([("python", "#"), ("rust", "//"),]);
 }