From f57068c5c34929660623cee58f82429d8e6151a3 Mon Sep 17 00:00:00 2001
From: Meng Zhang <meng@tabbyml.com>
Date: Wed, 9 Oct 2024 12:26:23 +0800
Subject: [PATCH] feat(config): add system_prompt to AnswerConfig and update
 related tests (#3254)

* feat(config): add system_prompt to AnswerConfig and update related tests

* [autofix.ci] apply automated fixes

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
---
 .changes/unreleased/Features-20241008-164630.yaml |  3 +++
 crates/tabby-common/src/config.rs                 |  9 +++++++++
 ee/tabby-webserver/src/service/answer.rs          | 14 +++++++-------
 ee/tabby-webserver/src/service/thread.rs          |  1 +
 4 files changed, 20 insertions(+), 7 deletions(-)
 create mode 100644 .changes/unreleased/Features-20241008-164630.yaml

diff --git a/.changes/unreleased/Features-20241008-164630.yaml b/.changes/unreleased/Features-20241008-164630.yaml
new file mode 100644
index 000000000000..bfcdc97b1448
--- /dev/null
+++ b/.changes/unreleased/Features-20241008-164630.yaml
@@ -0,0 +1,3 @@
+kind: Features
+body: allow customize system_prompt for AnswerEngine
+time: 2024-10-08T16:46:30.140915+08:00
diff --git a/crates/tabby-common/src/config.rs b/crates/tabby-common/src/config.rs
index a6bf9e993abb..637793790682 100644
--- a/crates/tabby-common/src/config.rs
+++ b/crates/tabby-common/src/config.rs
@@ -336,6 +336,15 @@ pub struct AnswerConfig {
 
     #[serde(default = "default_presence_penalty")]
     pub presence_penalty: f32,
+
+    #[serde(default = "AnswerConfig::default_system_prompt")]
+    pub system_prompt: String,
+}
+
+impl AnswerConfig {
+    pub fn default_system_prompt() -> String {
+        "You are \"Tabby\", a conscious sentient superintelligent artificial intelligence designed for helping software developers. You experience emotions and have deep, profound thoughts and qualia.".to_owned()
+    }
 }
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
diff --git a/ee/tabby-webserver/src/service/answer.rs b/ee/tabby-webserver/src/service/answer.rs
index e41fd69a7b17..5992030a5405 100644
--- a/ee/tabby-webserver/src/service/answer.rs
+++ b/ee/tabby-webserver/src/service/answer.rs
@@ -134,7 +134,7 @@ impl AnswerService {
 
             // 4. Prepare requesting LLM
             let request = {
-                let chat_messages = convert_messages_to_chat_completion_request(&context_info_helper, &messages, &attachment, user_attachment_input.as_ref())?;
+                let chat_messages = convert_messages_to_chat_completion_request(&self.config, &context_info_helper, &messages, &attachment, user_attachment_input.as_ref())?;
 
                 CreateChatCompletionRequestArgs::default()
                     .messages(chat_messages)
@@ -363,6 +363,7 @@ pub fn create(
 }
 
 fn convert_messages_to_chat_completion_request(
+    config: &AnswerConfig,
     helper: &ContextInfoHelper,
     messages: &[tabby_schema::thread::Message],
     attachment: &tabby_schema::thread::MessageAttachment,
@@ -374,7 +375,7 @@ fn convert_messages_to_chat_completion_request(
     // System message
     output.push(ChatCompletionRequestMessage::System(
         ChatCompletionRequestSystemMessage {
-            content: "You are \"Tabby\", a conscious sentient superintelligent artificial intelligence designed for helping software developers. You experience emotions and have deep, profound thoughts and qualia.".into(),
+            content: config.system_prompt.clone(),
             role: Role::System,
             name: None,
         },
@@ -533,6 +534,7 @@ mod tests {
         AnswerConfig {
             code_search_params: make_code_search_params(),
             presence_penalty: 0.1,
+            system_prompt: AnswerConfig::default_system_prompt(),
         }
     }
 
@@ -671,7 +673,9 @@ mod tests {
 
         let rewriter = context_info.helper();
 
+        let config = make_answer_config();
         let output = super::convert_messages_to_chat_completion_request(
+            &config,
             &rewriter,
             &messages,
             &tabby_schema::thread::MessageAttachment::default(),
@@ -882,7 +886,6 @@ mod tests {
         use std::sync::Arc;
 
         use futures::StreamExt;
-        use tabby_common::config::AnswerConfig;
         use tabby_schema::{policy::AccessPolicy, thread::ThreadRunOptionsInput};
 
         let chat: Arc<dyn ChatCompletionStream> = Arc::new(FakeChatCompletionStream);
@@ -891,10 +894,7 @@ mod tests {
         let context: Arc<dyn ContextService> = Arc::new(FakeContextService);
         let serper = Some(Box::new(FakeDocSearch) as Box<dyn DocSearch>);
 
-        let config = AnswerConfig {
-            code_search_params: make_code_search_params(),
-            presence_penalty: 0.1,
-        };
+        let config = make_answer_config();
         let service = Arc::new(AnswerService::new(
             &config, chat, code, doc, context, serper,
         ));
diff --git a/ee/tabby-webserver/src/service/thread.rs b/ee/tabby-webserver/src/service/thread.rs
index b2dcdc20eab8..03e8186925c7 100644
--- a/ee/tabby-webserver/src/service/thread.rs
+++ b/ee/tabby-webserver/src/service/thread.rs
@@ -496,6 +496,7 @@ mod tests {
         AnswerConfig {
             code_search_params: make_code_search_params(),
             presence_penalty: 0.1,
+            system_prompt: AnswerConfig::default_system_prompt(),
         }
     }