From 1cc89eff514669d6459f62cc4eed8e0b34d8ef0c Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Tue, 23 Apr 2024 14:32:06 +0800
Subject: refactor: more async code (#427)

---
 src/repl/mod.rs | 33 ++++++++++++++++-----------------
 1 file changed, 16 insertions(+), 17 deletions(-)

(limited to 'src/repl')

diff --git a/src/repl/mod.rs b/src/repl/mod.rs
index 412d0be..3c0b98b 100644
--- a/src/repl/mod.rs
+++ b/src/repl/mod.rs
@@ -6,9 +6,9 @@ use self::completer::ReplCompleter;
 use self::highlighter::ReplHighlighter;
 use self::prompt::ReplPrompt;
 
-use crate::client::{ensure_model_capabilities, init_client};
+use crate::client::{ensure_model_capabilities, init_client, send_stream};
 use crate::config::{GlobalConfig, Input, InputContext, State};
-use crate::render::{render_error, render_stream};
+use crate::render::render_error;
 use crate::utils::{create_abort_signal, set_text, AbortSignal};
 
 use anyhow::{bail, Context, Result};
@@ -93,7 +93,7 @@ impl Repl {
         })
     }
 
-    pub fn run(&mut self) -> Result<()> {
+    pub async fn run(&mut self) -> Result<()> {
         self.banner();
 
         loop {
@@ -104,7 +104,7 @@ impl Repl {
             match sig {
                 Ok(Signal::Success(line)) => {
                     self.abort.reset();
-                    match self.handle(&line) {
+                    match self.handle(&line).await {
                         Ok(exit) => {
                             if exit {
                                 break;
@@ -127,11 +127,11 @@ impl Repl {
                 _ => {}
             }
         }
-        self.handle(".exit session")?;
+        self.handle(".exit session").await?;
         Ok(())
     }
 
-    fn handle(&self, mut line: &str) -> Result<bool> {
+    async fn handle(&self, mut line: &str) -> Result<bool> {
         if let Ok(Some(captures)) = MULTILINE_RE.captures(line) {
             if let Some(text_match) = captures.get(1) {
                 line = text_match.as_str();
@@ -175,7 +175,7 @@ impl Repl {
                             let role = self.config.read().retrieve_role(name.trim())?;
                             let input =
                                 Input::from_str(text.trim(), InputContext::new(Some(role), false));
-                            self.ask(input)?;
+                            self.ask(input).await?;
                         }
                         None => {
                             self.config.write().set_role(args)?;
@@ -220,7 +220,7 @@ impl Repl {
                         };
                         let files = shell_words::split(files).with_context(|| "Invalid args")?;
                         let input = Input::new(text, files, self.config.read().input_context())?;
-                        self.ask(input)?;
+                        self.ask(input).await?;
                     }
                     None => println!("Usage: .file <files>... [-- <text>...]"),
                 },
@@ -246,7 +246,7 @@ impl Repl {
             },
             None => {
                 let input = Input::from_str(line, self.config.read().input_context());
-                self.ask(input)?;
+                self.ask(input).await?;
             }
         }
 
@@ -255,7 +255,7 @@ impl Repl {
         Ok(false)
     }
 
-    fn ask(&self, input: Input) -> Result<()> {
+    async fn ask(&self, input: Input) -> Result<()> {
         if input.is_empty() {
             return Ok(());
         }
@@ -265,7 +265,7 @@ impl Repl {
         self.config.read().maybe_print_send_tokens(&input);
         let mut client = init_client(&self.config)?;
         ensure_model_capabilities(client.as_mut(), input.required_capabilities())?;
-        let output = render_stream(&input, client.as_ref(), &self.config, self.abort.clone())?;
+        let output = send_stream(&input, client.as_ref(), &self.config, self.abort.clone()).await?;
         self.config.write().save_message(input, &output)?;
         self.config.read().maybe_copy(&output);
         if self.config.write().should_compress_session() {
@@ -283,10 +283,9 @@ impl Repl {
                 color.italic().paint("compress_threshold"),
                 color.normal().paint("`."),
             );
-            std::thread::spawn(move || -> anyhow::Result<()> {
-                let _ = compress_session(&config);
+            tokio::spawn(async move {
+                let _ = compress_session(&config).await;
                 config.write().end_compressing_session();
-                Ok(())
             });
         }
         Ok(())
@@ -443,14 +442,14 @@ fn parse_command(line: &str) -> Option<(&str, Option<&str>)> {
     }
 }
 
-fn compress_session(config: &GlobalConfig) -> Result<()> {
+async fn compress_session(config: &GlobalConfig) -> Result<()> {
     let input = Input::from_str(
-        &config.read().summarize_prompt,
+        config.read().summarize_prompt(),
         config.read().input_context(),
     );
     let mut client = init_client(config)?;
     ensure_model_capabilities(client.as_mut(), input.required_capabilities())?;
-    let summary = client.send_message(input)?;
+    let summary = client.send_message(input).await?;
     config.write().compress_session(&summary);
     Ok(())
 }
-- 
cgit v1.2.3