From ba884c9fc65d820f8603859acafd055b98897b56 Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Mon, 17 Jun 2024 13:19:00 +0800
Subject: feat: add `.continue` repl command (#608)

---
 src/config/session.rs | 40 +++++++++++++++++++++++++++-------------
 1 file changed, 27 insertions(+), 13 deletions(-)

(limited to 'src/config/session.rs')

diff --git a/src/config/session.rs b/src/config/session.rs
index b95fcc5..df9b082 100644
--- a/src/config/session.rs
+++ b/src/config/session.rs
@@ -353,20 +353,31 @@ impl Session {
     }
 
     pub fn add_message(&mut self, input: &Input, output: &str) -> Result<()> {
-        let mut need_add_msg = true;
-        if self.messages.is_empty() {
-            self.messages.extend(input.role().build_messages(input));
-            need_add_msg = false;
-        }
-        if need_add_msg {
-            self.messages
-                .push(Message::new(MessageRole::User, input.message_content()));
+        match input.continue_output() {
+            Some(_) => {
+                if let Some(message) = self.messages.last_mut() {
+                    if let MessageContent::Text(text) = &mut message.content {
+                        *text = format!("{text}{output}");
+                    }
+                }
+            }
+            None => {
+                let mut need_add_msg = true;
+                if self.messages.is_empty() {
+                    self.messages.extend(input.role().build_messages(input));
+                    need_add_msg = false;
+                }
+                if need_add_msg {
+                    self.messages
+                        .push(Message::new(MessageRole::User, input.message_content()));
+                }
+                self.data_urls.extend(input.data_urls());
+                self.messages.push(Message::new(
+                    MessageRole::Assistant,
+                    MessageContent::Text(output.to_string()),
+                ));
+            }
         }
-        self.data_urls.extend(input.data_urls());
-        self.messages.push(Message::new(
-            MessageRole::Assistant,
-            MessageContent::Text(output.to_string()),
-        ));
         self.dirty = true;
         Ok(())
     }
@@ -385,6 +396,9 @@ impl Session {
 
     pub fn build_messages(&self, input: &Input) -> Vec<Message> {
         let mut messages = self.messages.clone();
+        if input.continue_output().is_some() {
+            return messages;
+        }
         let mut need_add_msg = true;
         let len = messages.len();
         if len == 0 {
-- 
cgit v1.2.3