feat: add rust prompt rewrite support (#296)

Cargo.lock

@@ -2851,6 +2851,7 @@ dependencies = [
  "tracing-test",
  "tree-sitter-javascript",
  "tree-sitter-python",
+ "tree-sitter-rust",
  "tree-sitter-tags",
  "walkdir",
 ]
@@ -3528,6 +3529,16 @@ dependencies = [
  "tree-sitter",
 ]
 
+[[package]]
+name = "tree-sitter-rust"
+version = "0.20.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "797842733e252dc11ae5d403a18060bf337b822fc2ae5ddfaa6ff4d9cc20bda6"
+dependencies = [
+ "cc",
+ "tree-sitter",
+]
+
 [[package]]
 name = "tree-sitter-tags"
 version = "0.20.2"

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"

crates/tabby-scheduler/src/dataset.rs

@@ -224,7 +224,8 @@ lazy_static! {
         map
     };
     static ref LANGUAGE_TAGS: HashMap<&'static str, TagsConfigurationSync> = {
-        HashMap::from([(
+        HashMap::from([
+            (
                 "python",
                 TagsConfigurationSync(
                     TagsConfiguration::new(
@@ -234,6 +235,18 @@ lazy_static! {
                     )
                     .unwrap(),
                 ),
-        )])
+            ),
+            (
+                "rust",
+                TagsConfigurationSync(
+                    TagsConfiguration::new(
+                        tree_sitter_rust::language(),
+                        tree_sitter_rust::TAGGING_QUERY,
+                        "",
+                    )
+                    .unwrap(),
+                ),
+            ),
+        ])
     };
 }

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", "//"),]);
 }