From 598c31c3d946e8772c4d1bc450b5d2b1bdf587c0 Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Sun, 5 Mar 2023 07:52:36 +0800
Subject: refactor: optimize role prompt (#13)

---
 src/client.rs | 32 ++++++++++++++++++++++----------
 src/repl.rs   |  8 ++++----
 2 files changed, 26 insertions(+), 14 deletions(-)

(limited to 'src')

diff --git a/src/client.rs b/src/client.rs
index 1f42346..3fe75dd 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -81,11 +81,10 @@ impl ChatGptClient {
     }
 
     async fn acquire_inner(&self, content: &str, prompt: Option<String>) -> Result<String> {
-        let content = combine(content, prompt);
         if self.config.dry_run {
-            return Ok(content);
+            return Ok(combine(content, prompt));
         }
-        let builder = self.request_builder(&content, false);
+        let builder = self.request_builder(content, prompt, false);
 
         let data: Value = builder.send().await?.json().await?;
 
@@ -102,12 +101,11 @@ impl ChatGptClient {
         prompt: Option<String>,
         receiver: &mut ReplyReceiver,
     ) -> Result<()> {
-        let content = combine(content, prompt);
         if self.config.dry_run {
-            receiver.text(&content);
+            receiver.text(&combine(content, prompt));
             return Ok(());
         }
-        let builder = self.request_builder(&content, true);
+        let builder = self.request_builder(content, prompt, true);
         let mut stream = builder.send().await?.bytes_stream().eventsource();
         let mut virgin = true;
         while let Some(part) = stream.next().await {
@@ -135,10 +133,25 @@ impl ChatGptClient {
         Ok(())
     }
 
-    fn request_builder(&self, content: &str, stream: bool) -> RequestBuilder {
+    fn request_builder(
+        &self,
+        content: &str,
+        prompt: Option<String>,
+        stream: bool,
+    ) -> RequestBuilder {
+        let user_message = json!({ "role": "user", "content": content });
+        let messages = match prompt {
+            Some(prompt) => {
+                let system_message = json!({ "role": "system", "content": prompt.trim() });
+                json!([system_message, user_message])
+            }
+            None => {
+                json!([user_message])
+            }
+        };
         let mut body = json!({
             "model": MODEL,
-            "messages": [{"role": "user", "content": content}],
+            "messages": messages,
         });
 
         if let Some(v) = self.config.temperature {
@@ -160,11 +173,10 @@ impl ChatGptClient {
 
 fn combine(content: &str, prompt: Option<String>) -> String {
     match prompt {
-        Some(v) => format!("{v} {content}"),
+        Some(prompt) => format!("{}\n{content}", prompt.trim()),
         None => content.to_string(),
     }
 }
-
 fn init_runtime() -> Result<Runtime> {
     tokio::runtime::Builder::new_current_thread()
         .enable_all()
diff --git a/src/repl.rs b/src/repl.rs
index 164d746..17617f8 100644
--- a/src/repl.rs
+++ b/src/repl.rs
@@ -134,7 +134,6 @@ impl Repl {
                 },
                 ".clear-role" => {
                     handler.handle(ReplCmd::UnsetRole)?;
-                    dump("", 1);
                 }
                 ".editor" => {
                     dump(
@@ -156,7 +155,6 @@ impl Repl {
                 }
                 ".info" => {
                     handler.handle(ReplCmd::Info)?;
-                    dump("", 1);
                 }
                 _ => dump_unknown_command(),
             }
@@ -295,8 +293,9 @@ impl ReplCmdHandler {
             }
             ReplCmd::SetRole(name) => match self.config.find_role(&name) {
                 Some(role) => {
+                    let output = format!("{}>> {}", role.name, role.prompt.trim());
                     self.state.borrow_mut().role = Some(role);
-                    dump("", 1);
+                    dump(output, 2);
                 }
                 None => {
                     dump("Unknown role", 2);
@@ -304,6 +303,7 @@ impl ReplCmdHandler {
             },
             ReplCmd::UnsetRole => {
                 self.state.borrow_mut().role = None;
+                dump("Clear the current selected role", 1);
             }
             ReplCmd::Info => {
                 let state = self.state.borrow();
@@ -338,7 +338,7 @@ impl ReplCmdHandler {
                 for (name, value) in items {
                     info.push_str(&format!("{name:<20}{value}\n"));
                 }
-                dump(info, 0);
+                dump(info, 1);
             }
         }
         Ok(())
-- 
cgit v1.2.3