diff --git a/crates/tabby/src/serve/completions/prompt.rs b/crates/tabby/src/serve/completions/prompt.rs index cf98727..5b94f7d 100644 --- a/crates/tabby/src/serve/completions/prompt.rs +++ b/crates/tabby/src/serve/completions/prompt.rs @@ -49,11 +49,7 @@ impl PromptBuilder { pub fn build(&self, language: &str, segments: Segments) -> String { let segments = self.rewrite(language, segments); - if let Some(suffix) = segments.suffix { - self.build_prompt(segments.prefix, suffix) - } else { - self.build_prompt(segments.prefix, "\n".to_owned()) - } + self.build_prompt(segments.prefix, get_default_suffix(segments.suffix)) } fn rewrite(&self, language: &str, segments: Segments) -> Segments { @@ -65,6 +61,19 @@ impl PromptBuilder { } } +fn get_default_suffix(suffix: Option) -> String { + if suffix.is_none() { + return "\n".to_owned(); + } + + let suffix = suffix.unwrap(); + if suffix.is_empty() { + "\n".to_owned() + } else { + suffix + } +} + fn rewrite_with_index(index: &IndexState, language: &str, segments: Segments) -> Segments { let snippets = collect_snippets(index, language, &segments.prefix); if snippets.is_empty() { @@ -261,7 +270,7 @@ mod tests { }; assert_eq!( pb.build(language, segments), - "
 this is some prefix  "
+                "
 this is some prefix \n "
             );
         }
 
@@ -292,7 +301,7 @@ mod tests {
                 prefix: "".into(),
                 suffix: Some("".into()),
             };
-            assert_eq!(pb.build(language, segments), "
   ");
+            assert_eq!(pb.build(language, segments), "
  \n ");
         }
     }