From 775576b53eddae86bb3db68fcaf13241271380b3 Mon Sep 17 00:00:00 2001
From: Meng Zhang <meng@tabbyml.com>
Date: Sat, 3 Jun 2023 10:29:04 -0700
Subject: [PATCH] fix: only use prompt_template when suffix presents [TAB-46]
 (#184)

* fix: only use prompt_template when suffix presents

* lint
---
 crates/tabby/src/download/cache_info.rs | 14 --------------
 crates/tabby/src/serve/completions.rs   | 11 ++++++++---
 2 files changed, 8 insertions(+), 17 deletions(-)

diff --git a/crates/tabby/src/download/cache_info.rs b/crates/tabby/src/download/cache_info.rs
index d746b9e..59972c2 100644
--- a/crates/tabby/src/download/cache_info.rs
+++ b/crates/tabby/src/download/cache_info.rs
@@ -53,17 +53,3 @@ impl CacheInfo {
         Ok(())
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[tokio::test]
-    async fn test_hf() {
-        let hf_metadata = HFCacheInfo::from("TabbyML/J-350M").await;
-        assert_eq!(
-            hf_metadata.transformers_info.auto_model,
-            "AutoModelForCausalLM"
-        );
-    }
-}
diff --git a/crates/tabby/src/serve/completions.rs b/crates/tabby/src/serve/completions.rs
index 2c5b5c2..4ed567f 100644
--- a/crates/tabby/src/serve/completions.rs
+++ b/crates/tabby/src/serve/completions.rs
@@ -33,7 +33,7 @@ pub struct Segments {
 
     /// Content that appears after the cursor in the editor window.
     #[schema(example = "\n        return fib(n - 1) + fib(n - 2)")]
-    suffix: String,
+    suffix: Option<String>,
 }
 
 #[derive(Serialize, Deserialize, ToSchema, Clone, Debug)]
@@ -65,8 +65,13 @@ pub async fn completion(
 
     let prompt = if let Some(Segments { prefix, suffix }) = request.segments {
         if let Some(prompt_template) = &state.prompt_template {
-            strfmt!(prompt_template, prefix => prefix, suffix => suffix)
-                .expect("Failed to format prompt")
+            if let Some(suffix) = suffix {
+                strfmt!(prompt_template, prefix => prefix, suffix => suffix)
+                    .expect("Failed to format prompt")
+            } else {
+                // If suffix is empty, just returns prefix.
+                prefix
+            }
         } else {
             // If there's no prompt template, just use prefix.
             prefix