diff options
| author | sigoden <sigoden@gmail.com> | 2024-06-05 09:02:23 +0800 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2024-06-05 09:02:23 +0800 |
| commit | 1ec6abfaee2fdc189b348b7e3a8145bd9a84da74 (patch) | |
| tree | 6f68a860a39b0fbe87784de925b5a31e00e74e33 /src/utils/spinner.rs | |
| parent | 71f2e94579511d7524f5534377001ab3f02a9597 (diff) | |
| download | aichat-1ec6abfaee2fdc189b348b7e3a8145bd9a84da74.tar.gz | |
feat: support RAG (#560)
* feat: support RAG
* support more embeddings models and implement concurrent embedding api
* show the progress of addings paths
* ignore embedding context when saving message
* embedding model max_chunk_size => default_chunk_size
* support pdf and pandoc formats (docx, epub, ipynb)
Diffstat (limited to 'src/utils/spinner.rs')
| -rw-r--r-- | src/utils/spinner.rs | 39 |
1 files changed, 33 insertions, 6 deletions
diff --git a/src/utils/spinner.rs b/src/utils/spinner.rs index 0dd4d01..c746469 100644 --- a/src/utils/spinner.rs +++ b/src/utils/spinner.rs @@ -4,7 +4,10 @@ use std::{ io::{stdout, Stdout, Write}, time::Duration, }; -use tokio::{sync::oneshot, time::interval}; +use tokio::{ + sync::{mpsc, oneshot}, + time::interval, +}; pub struct Spinner { index: usize, @@ -23,6 +26,10 @@ impl Spinner { } } + pub fn set_message(&mut self, message: &str) { + self.message = format!(" {message}"); + } + pub fn step(&mut self, writer: &mut Stdout) -> Result<()> { if self.stopped { return Ok(()); @@ -55,18 +62,38 @@ impl Spinner { } } -pub async fn run_spinner(message: &str, rx: oneshot::Receiver<()>) -> Result<()> { +pub async fn run_spinner(message: &str) -> (oneshot::Sender<()>, mpsc::UnboundedSender<String>) { + let message = format!(" {message}"); + let (stop_tx, stop_rx) = oneshot::channel(); + let (message_tx, message_rx) = mpsc::unbounded_channel(); + tokio::spawn(run_spinner_inner(message, stop_rx, message_rx)); + (stop_tx, message_tx) +} + +async fn run_spinner_inner( + message: String, + stop_rx: oneshot::Receiver<()>, + mut message_rx: mpsc::UnboundedReceiver<String>, +) -> Result<()> { let mut writer = stdout(); - let mut spinner = Spinner::new(message); + let mut spinner = Spinner::new(&message); let mut interval = interval(Duration::from_millis(50)); tokio::select! { _ = async { loop { - interval.tick().await; - let _ = spinner.step(&mut writer); + tokio::select! { + _ = interval.tick() => { + let _ = spinner.step(&mut writer); + } + message = message_rx.recv() => { + if let Some(message) = message { + spinner.set_message(&message); + } + } + } } } => {} - _ = rx => { + _ = stop_rx => { spinner.stop(&mut writer)?; } } |
