refactor(openai): support fim and no_fim for openai completion

crates/http-api-bindings/src/completion/mistral.rs

@@ -5,7 +5,7 @@
 use serde::{Deserialize, Serialize};
 use tabby_inference::{CompletionOptions, CompletionStream};
 
-use super::FIM_TOKEN;
+use super::split_fim_prompt;
 
 pub struct MistralFIMEngine {
     client: reqwest::Client,
@@ -68,13 +68,10 @@
 #[async_trait]
 impl CompletionStream for MistralFIMEngine {
     async fn generate(&self, prompt: &str, options: CompletionOptions) -> BoxStream<String> {
-        let parts = prompt.splitn(2, FIM_TOKEN).collect::<Vec<_>>();
+        let (prompt, suffix) = split_fim_prompt(prompt);
         let request = FIMRequest {
-            prompt: parts[0].to_owned(),
-            suffix: parts
-                .get(1)
-                .map(|x| x.to_string())
-                .filter(|x| !x.is_empty()),
+            prompt: prompt.to_owned(),
+            suffix: suffix.map(|x| x.to_owned()),
             model: self.model_name.clone(),
             max_tokens: options.max_decoding_tokens,
             temperature: options.sampling_temperature,

crates/http-api-bindings/src/completion/mod.rs

@@ -31,14 +31,27 @@
             );
             Arc::new(engine)
         }
-        "openai/completion" => {
+        "openai/legacy_completion" | "openai/completion" | "deepseek/completion" => {
             let engine = OpenAICompletionEngine::create(
                 model.model_name.clone(),
                 model
                     .api_endpoint
                     .as_deref()
                     .expect("api_endpoint is required"),
                 model.api_key.clone(),
+                true,
+            );
+            Arc::new(engine)
+        }
+        "openai/legacy_completion_no_fim" | "vllm/completion" => {
+            let engine = OpenAICompletionEngine::create(
+                model.model_name.clone(),
+                model
+                    .api_endpoint
+                    .as_deref()
+                    .expect("api_endpoint is required"),
+                model.api_key.clone(),
+                false,
             );
             Arc::new(engine)
         }
@@ -59,3 +72,29 @@
         (model.prompt_template.clone(), model.chat_template.clone())
     }
 }
+
+fn split_fim_prompt(prompt: &str) -> (&str, Option<&str>) {
+    let parts = prompt.splitn(2, FIM_TOKEN).collect::<Vec<_>>();
+    (parts[0], parts.get(1).copied())
+}
+
+#[cfg(test)]
+mod tests {
+    use std::vec;
+
+    use super::*;
+
+    #[test]
+    fn test_split_fim_prompt_no_fim() {
+        let no_fim = vec![
+            ("prefix<|FIM|>suffix", ("prefix", Some("suffix"))),
+            ("prefix<|FIM|>", ("prefix", Some(""))),
+            ("<|FIM|>suffix", ("", Some("suffix"))),
+            ("<|FIM|>", ("", Some(""))),
+            ("prefix", ("prefix", None)),
+        ];
+        for (input, expected) in no_fim {
+            assert_eq!(split_fim_prompt(input), expected);
+        }
+    }
+}

crates/http-api-bindings/src/completion/openai.rs

@@ -5,17 +5,27 @@
 use serde::{Deserialize, Serialize};
 use tabby_inference::{CompletionOptions, CompletionStream};
 
-use super::FIM_TOKEN;
+use super::split_fim_prompt;
 
 pub struct OpenAICompletionEngine {
     client: reqwest::Client,
     model_name: String,
     api_endpoint: String,
     api_key: Option<String>,
+
+    /// OpenAI Completion API use suffix field in request when FIM is not supported,
+    /// support_fim is used to mark if FIM is supported,
+    /// provide a `openai/legacy_completion_no_fim` backend to use suffix field.
+    support_fim: bool,
 }
 
 impl OpenAICompletionEngine {
-    pub fn create(model_name: Option<String>, api_endpoint: &str, api_key: Option<String>) -> Self {
+    pub fn create(
+        model_name: Option<String>,
+        api_endpoint: &str,
+        api_key: Option<String>,
+        support_fim: bool,
+    ) -> Self {
         let model_name = model_name.expect("model_name is required for openai/completion");
         let client = reqwest::Client::new();
 
@@ -24,6 +34,7 @@
             model_name,
             api_endpoint: format!("{}/completions", api_endpoint),
             api_key,
+            support_fim,
         }
     }
 }
@@ -53,14 +64,16 @@
 #[async_trait]
 impl CompletionStream for OpenAICompletionEngine {
     async fn generate(&self, prompt: &str, options: CompletionOptions) -> BoxStream<String> {
-        let parts = prompt.splitn(2, FIM_TOKEN).collect::<Vec<_>>();
+        let (prompt, suffix) = if self.support_fim {
+            split_fim_prompt(prompt)
+        } else {
+            (prompt, None)
+        };
+
         let request = CompletionRequest {
             model: self.model_name.clone(),
-            prompt: parts[0].to_owned(),
-            suffix: parts
-                .get(1)
-                .map(|x| x.to_string())
-                .filter(|x| !x.is_empty()),
+            prompt: prompt.to_owned(),
+            suffix: suffix.map(|x| x.to_owned()),
             max_tokens: options.max_decoding_tokens,
             temperature: options.sampling_temperature,
             stream: true,