fix: ensure default suffix to be non-empty (#400)

release-0.2
Meng Zhang 2023-09-05 11:45:29 +08:00 committed by GitHub
parent abfa7975e8
commit d85cd81139
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 16 additions and 7 deletions

View File

@ -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>) -> 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),
"<PRE> this is some prefix <SUF> <MID>"
"<PRE> this is some prefix <SUF>\n <MID>"
);
}
@ -292,7 +301,7 @@ mod tests {
prefix: "".into(),
suffix: Some("".into()),
};
assert_eq!(pb.build(language, segments), "<PRE> <SUF> <MID>");
assert_eq!(pb.build(language, segments), "<PRE> <SUF>\n <MID>");
}
}